diff --git a/.gitignore b/.gitignore index 256048b..160afdb 100755 --- a/.gitignore +++ b/.gitignore @@ -138,6 +138,9 @@ docs/**/modelzoo.md *.mp4 *.DS_Store *.png +!tests/assets/ +!tests/assets/ltx25/ +!tests/assets/ltx25/official_guitar_man.png *.csv *.json *.jpg diff --git a/CLAUDE.md b/CLAUDE.md index 0b093ae..7761db2 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -102,6 +102,8 @@ When adding or porting a pipeline: - Select the closest maintained pipeline, public example, and tests as structural baselines. Read the relevant adding-new-example, adding-new-model, adding-new-stage, model-loading, configuration, and service guides. +- Base each new model-family example README on `examples/README_TEMPLATE.md`; keep its required section order and + remove inapplicable optional sections and all placeholders. - Inventory model-specific classes and configuration fields, then map them to upstream behavior and the selected baseline. - Reuse `BasePipeline`, `BaseStage`, `ModuleManager`, existing configuration dataclasses, example contracts, and diff --git a/README.md b/README.md index d80d163..6b376d5 100644 --- a/README.md +++ b/README.md @@ -25,6 +25,9 @@ runtime path, supported workloads, and reproducible real-time gate. ## News 📰 +- ✨ **2026-08-19**: Added [**LTX-2.5 Distilled**](examples/ltx25_distilled/README.md) T2V and I2V joint + audio-video generation with a ModuleManager-backed six-stage pipeline, selectable dense attention backends, and + Ulysses sequence parallelism on **1, 2, or 4 x H100** GPUs. - ✨ **2026-08-10**: Added [**SwiftVR**](examples/swiftvr/README.md) causal video restoration with stateful streaming and a single-GPU H100 example. - ✨ **2026-08-05**: Added [**MiniMax H3**](examples/minimax_h3/README.md) T2VA, FL2VA, and Ref2VA joint @@ -247,6 +250,7 @@ telefuser/ |----------|------|-------| | `WanVideo` (Wan2.1 / Wan2.2) | T2V, I2V, FL2V | Main video generation family, including async and service examples in [examples/wan_video/README.md](examples/wan_video/README.md) | | `LTX Video` | I2V + Audio | Unified audio-video generation via [examples/ltx_video/README.md](examples/ltx_video/README.md) | +| `LTX-2.5 Distilled` | T2V, I2V + Audio | ModuleManager-backed six-stage pipeline with 1/2/4-H100 Ulysses SP; see [examples/ltx25_distilled/README.md](examples/ltx25_distilled/README.md) | | `MiniMax H3` | T2VA, FL2VA, Ref2VA + Audio | Local 768p joint audio-video generation via [examples/minimax_h3/README.md](examples/minimax_h3/README.md) | | `LongCat-Video` | T2V, I2V, VC | Long-form generation and continuation via [examples/longcat_video/README.md](examples/longcat_video/README.md) | | `LingBot-Video` | T2I, T2V, TI2V, MoE refiner | Dense/MoE generation with native CFG/SP and an in-memory base-to-refiner path; see [examples/lingbot_video/README.md](examples/lingbot_video/README.md) | diff --git a/README_zh.md b/README_zh.md index b5aff3e..42be36b 100644 --- a/README_zh.md +++ b/README_zh.md @@ -24,6 +24,9 @@ TeleFuser 是一个开源的多模态生成与世界模型流式推理和服务 ## News 📰 +- ✨ **2026-08-19**：新增 [**LTX-2.5 Distilled**](examples/ltx25_distilled/README.md) T2V 和 I2V 联合 + 音视频生成，采用基于 ModuleManager 的六阶段 Pipeline，支持选择密集注意力后端，并可在 + **1、2 或 4 张 H100** 上使用 Ulysses 序列并行。 - ✨ **2026-08-10**：新增 [**SwiftVR**](examples/swiftvr/README.md) 因果视频修复、有状态流式推理及单卡 H100 示例。 - ✨ **2026-08-05**：新增 [**MiniMax H3**](examples/minimax_h3/README.md) T2VA、FL2VA 和 Ref2VA 联合 音视频生成，并支持标准 `telefuser serve` 服务模式。在相同的 768p、5 秒、50 步 T2VA 请求和一次预热 @@ -233,6 +236,7 @@ telefuser/ |----------|------|------| | `WanVideo` (Wan2.1 / Wan2.2) | T2V, I2V, FL2V | 主力视频生成家族，含异步和服务示例，见 [examples/wan_video/README.md](examples/wan_video/README.md) | | `LTX Video` | I2V + Audio | 统一音视频生成，见 [examples/ltx_video/README.md](examples/ltx_video/README.md) | +| `LTX-2.5 Distilled` | T2V, I2V + Audio | 基于 ModuleManager 的六阶段 Pipeline，支持 1/2/4 张 H100 的 Ulysses SP，见 [examples/ltx25_distilled/README.md](examples/ltx25_distilled/README.md) | | `MiniMax H3` | T2VA, FL2VA, Ref2VA + Audio | 本地 768p 联合音视频生成，见 [examples/minimax_h3/README.md](examples/minimax_h3/README.md) | | `LongCat-Video` | T2V, I2V, VC | 长视频生成与续写，见 [examples/longcat_video/README.md](examples/longcat_video/README.md) | | **NEW** `LingBot-Video` | T2I, T2V, TI2V, MoE refiner | 支持原生 CFG/SP 的 Dense/MoE 生成与内存直传 base-to-refiner，见 [examples/lingbot_video/README.md](examples/lingbot_video/README.md) | diff --git a/docs/en/adding_new_example.md b/docs/en/adding_new_example.md index 233588b..f533bf7 100644 --- a/docs/en/adding_new_example.md +++ b/docs/en/adding_new_example.md @@ -11,6 +11,10 @@ Pipeline examples are standalone Python scripts that demonstrate how to use Tele 3. Compatible with the TeleFuser server (`telefuser serve`) 4. Well-documented with clear naming conventions +Each model-family directory must include a `README.md` based on +[`examples/README_TEMPLATE.md`](../../examples/README_TEMPLATE.md). Keep the required section order, remove unused +optional sections, and replace all template placeholders before submitting the example. + ## File Structure and Naming ### Directory Organization diff --git a/docs/en/index.md b/docs/en/index.md index 1a7d402..95cb3b2 100644 --- a/docs/en/index.md +++ b/docs/en/index.md @@ -73,6 +73,7 @@ Reusable stages, model configs, schedulers, and pipeline orchestration. |-------|-------|-------------| | WanVideo (Wan2.1 / Wan2.2) | T2V, I2V, FL2V | Video generation and editing | | LTX Video | I2V + Audio | Video generation with audio | +| LTX-2.5 Distilled | T2V, I2V + Audio | ModuleManager-backed six-stage pipeline with 1/2/4-H100 Ulysses SP; see the [example guide](../../examples/ltx25_distilled/README.md) | | MiniMax H3 | T2VA, FL2VA, Ref2VA + Audio | Local 768p joint audio-video generation | | FlashVSR | VSR | Video super-resolution | | SwiftVR | Causal video restoration | Stateful restoration with BF16, compile, FP8Linear, Ulysses SP, and stage-parallel options; see the [example guide](../../examples/swiftvr/README.md) | diff --git a/docs/zh/index.md b/docs/zh/index.md index 46c09c8..cddf53c 100644 --- a/docs/zh/index.md +++ b/docs/zh/index.md @@ -71,6 +71,7 @@ AdaTaylorCache 和运行时缓存控制，面向重复生成工作负载。 |------|------|------| | WanVideo (Wan2.1 / Wan2.2) | T2V, I2V, FL2V | 视频生成和编辑 | | LTX Video | I2V + Audio | 视频生成 + 音频 | +| LTX-2.5 Distilled | T2V、I2V + Audio | 基于 ModuleManager 的六阶段流水线，支持 1/2/4 张 H100 的 Ulysses SP，见[示例文档](../../examples/ltx25_distilled/README.md) | | MiniMax H3 | T2VA, FL2VA, Ref2VA + Audio | 本地 768p 音视频联合生成 | | FlashVSR | VSR | 视频超分辨率 | | SwiftVR | 因果视频修复 | 支持 BF16、torch.compile、FP8Linear、Ulysses SP 和 stage-parallel，见 [示例文档](../../examples/swiftvr/README.md) | diff --git a/examples/README_TEMPLATE.md b/examples/README_TEMPLATE.md new file mode 100644 index 0000000..06ec341 --- /dev/null +++ b/examples/README_TEMPLATE.md @@ -0,0 +1,192 @@ + + +# {MODEL_FAMILY} Examples + +{ONE_OR_TWO_SENTENCES_DESCRIBING_THE_SUPPORTED_MODELS_TASKS_AND_OUTPUTS} + +## Model Source + + + +| Model | HuggingFace | ModelScope | Purpose | +| --- | --- | --- | --- | +| `{MODEL_NAME}` | [{HF_REPOSITORY}]({HF_URL}) | [{MODELSCOPE_REPOSITORY}]({MODELSCOPE_URL}) | {MODEL_PURPOSE} | + +## Feature Support + + + +| Feature | Support | Notes | +| --- | --- | --- | +| {TASK_OR_FEATURE} | Supported | {CONSTRAINTS_OR_VARIANTS} | +| Multi-GPU inference | {SUPPORT_STATUS} | {PARALLEL_STRATEGY_AND_VALID_DEGREES} | +| LoRA | {SUPPORT_STATUS} | {SUPPORTED_VARIANTS} | +| Quantization | {SUPPORT_STATUS} | {DTYPES_OR_FORMATS} | +| CPU offload | {SUPPORT_STATUS} | {OFFLOAD_MODES} | +| Feature cache | {SUPPORT_STATUS} | {CACHE_IMPLEMENTATION} | +| Server API | {SUPPORT_STATUS} | {SERVE_OR_STREAM_SERVE} | + +## Requirements + + + +- GPU: {GPU_MODEL_OR_MINIMUM_VRAM} +- Software: {CUDA_PYTORCH_OR_EXTRA_PACKAGE_REQUIREMENTS} +- Input assets: {REQUIRED_INPUT_FORMATS_OR_NONE} + +Install TeleFuser by following the [development setup](../../CONTRIBUTING.md#development-setup). Then install any +example-specific dependencies: + +```bash +{INSTALL_COMMANDS_OR_COMMENT_STATING_NO_EXTRA_DEPENDENCIES} +``` + +## Model Directory + + + +```text +{MODEL_ROOT}/ +|-- {CHECKPOINT_OR_DIRECTORY} +\-- {AUXILIARY_CHECKPOINT_OR_DIRECTORY} +``` + +Set the model root if the examples use `TF_MODEL_ZOO_PATH`: + +```bash +export TF_MODEL_ZOO_PATH=/path/to/model_zoo +``` + +## Quick Start + + + +```bash +python examples/{example_directory}/{representative_script}.py \ + --model_root /path/to/model \ + --prompt "{EXAMPLE_PROMPT}" \ + --output_path work_dirs/{OUTPUT_FILE} +``` + +The command writes {OUTPUT_DESCRIPTION} to `work_dirs/{OUTPUT_FILE}`. + +## Examples + + + +### {TASK_NAME} + +#### `{script_name.py}` + +{ONE_SENTENCE_PURPOSE_AND_WHEN_TO_USE_THIS_SCRIPT} + +```bash +# Basic usage +python examples/{example_directory}/{script_name.py} \ + --model_root /path/to/model \ + {REQUIRED_ARGUMENTS} + +# Multi-GPU or another important variant +python examples/{example_directory}/{script_name.py} \ + --gpu_num {GPU_COUNT} \ + --model_root /path/to/model \ + {VARIANT_ARGUMENTS} +``` + +Key options: + +| Option | Default | Description | +| --- | --- | --- | +| `--model_root` | `{DEFAULT_OR_NONE}` | {MODEL_ROOT_DESCRIPTION} | +| `--gpu_num` | `{DEFAULT_GPU_COUNT}` | {GPU_COUNT_CONSTRAINTS} | +| `{OPTION}` | `{DEFAULT}` | {OPTION_DESCRIPTION} | + +Key behavior: + +- {IMPORTANT_DEFAULT_OR_MODEL_VARIANT} +- {OUTPUT_SHAPE_FORMAT_OR_LOCATION} +- {LIMITATION_OR_RESOURCE_NOTE} + +## Configuration + + + +### {CONFIGURATION_TOPIC} + +{EXPLAIN_THE_RULE_ITS_DEFAULT_AND_WHEN_TO_CHANGE_IT} + +```python +{MINIMAL_CONFIGURATION_SNIPPET} +``` + +## Serving + + + +Start the service: + +```bash +telefuser {serve_or_stream-serve} examples/{example_directory}/{server_script}.py \ + --port {PORT} \ + {OTHER_REQUIRED_OPTIONS} +``` + +See the [service guide](../../docs/en/service.md) or +[stream server guide](../../docs/en/stream_server.md) for API and deployment details. + +## Performance + + + +Measured with {GPU_COUNT_AND_MODEL}, {SOFTWARE_VERSIONS}, and commit `{GIT_REVISION}`. Results use {PRECISION}, +{ATTENTION_BACKEND}, and exclude {EXCLUDED_PHASES_OR_NOTHING}. + +| Configuration | GPUs | Resolution | Frames | Steps | Time (s) | Peak VRAM (GiB) | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| {CONFIGURATION_NAME} | {GPU_COUNT} | {RESOLUTION} | {FRAME_COUNT} | {STEP_COUNT} | {ELAPSED_TIME} | {PEAK_VRAM} | + +Reproduce the measurement: + +```bash +{BENCHMARK_COMMAND} +``` + +## Troubleshooting + + + +### {ERROR_OR_SYMPTOM} + +{CAUSE_AND_ACTIONABLE_FIX} + +```bash +{DIAGNOSTIC_OR_FIX_COMMAND} +``` + +## Notes + + + +- {MODEL_SPECIFIC_LIMITATION_OR_COMPATIBILITY_NOTE} +- {OUTPUT_OR_QUALITY_NOTE} diff --git a/examples/data/ltx25/README.md b/examples/data/ltx25/README.md new file mode 100644 index 0000000..0fc947b --- /dev/null +++ b/examples/data/ltx25/README.md @@ -0,0 +1,9 @@ +# LTX-2.5 I2V Reference Input + +`official_guitar_man.png` is the image referenced by the Lightricks LTX-2 model card's image-to-video example: + +https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/guitar-man.png + +Its SHA-256 is `e31cbbe4822ce07e1548121b436c0db3a067d1d78f2e75ab3e69375377b57274`. The matching official +prompt is `A man with short gray hair plays a red electric guitar.` Formal LTX-2.5 I2V regressions use this source at +frame 0 with strength 1.0. diff --git a/examples/data/ltx25/official_guitar_man.png b/examples/data/ltx25/official_guitar_man.png new file mode 100644 index 0000000..83ebe7d Binary files /dev/null and b/examples/data/ltx25/official_guitar_man.png differ diff --git a/examples/example_config.yaml b/examples/example_config.yaml index e6e3e1b..0b67991 100644 --- a/examples/example_config.yaml +++ b/examples/example_config.yaml @@ -101,6 +101,37 @@ pipelines: ppl_config_overrides: attn_impl: TORCH_SDPA + # ============================================================ + # ltx25_distilled - Joint Audio-Video Generation + # ============================================================ + + ltx25_distilled_t2v_2gpu: + script: ltx25_distilled/ltx25_distilled_t2v_h100.py + gpu_count: 2 + output_type: video + timeout_seconds: 3600 + height: 1024 + width: 1536 + use_run_with_file: true + require_audio: true + ppl_config_overrides: + num_frames: 121 + frame_rate: 24.0 + + ltx25_distilled_i2v_2gpu: + script: ltx25_distilled/ltx25_distilled_i2v_h100.py + gpu_count: 2 + output_type: video + timeout_seconds: 3600 + height: 512 + width: 896 + input_image_path: examples/data/ltx25/official_guitar_man.png + use_run_with_file: true + require_audio: true + ppl_config_overrides: + num_frames: 121 + frame_rate: 24.0 + # ============================================================ # minimax_h3 — Joint Audio-Video Generation # ============================================================ diff --git a/examples/ltx25_distilled/README.md b/examples/ltx25_distilled/README.md new file mode 100644 index 0000000..76f2f30 --- /dev/null +++ b/examples/ltx25_distilled/README.md @@ -0,0 +1,285 @@ +# LTX-2.5 Distilled Examples + +One-to-four-H100 text-to-video (T2V) and image-to-video (I2V) generation with the distilled LTX-2.5 pipeline. The example +produces an MP4 containing generated video and synchronized 48 kHz stereo audio. + +The LTX-2.5 implementation is isolated under `telefuser/models/ltx25` and +`telefuser/pipelines/ltx25_distilled`; it does not reuse the legacy LTX-2.3 model or pipeline modules. + +## Model Source + +| Model | HuggingFace | ModelScope | Purpose | +| --- | --- | --- | --- | +| LTX-2.5 22B distilled model pack | [Lightricks/LTX-2.5](https://huggingface.co/Lightricks/LTX-2.5) | N/A | Transformer, text encoder, video/audio VAEs, spatial upsampler, and duration head | + +This example does not auto-download weights. Download the official repository while preserving its directory layout: + +```bash +hf download Lightricks/LTX-2.5 --local-dir /path/to/LTX-2.5 +``` + +The example requires the exact split LTX-2.5 checkpoint layout shown below; a consolidated checkpoint or Diffusers +directory is not accepted. + +## Feature Support + +| Feature | Support | Notes | +| --- | --- | --- | +| Text-to-video | Supported | Generates video and audio from a text prompt | +| Image-to-video | Supported | Accepts one still image at a non-negative output-frame index | +| Multi-GPU inference | Supported | Ulysses sequence parallelism and block-level FSDP2 on 2 or 4 H100s | +| Attention backend | Supported | Select a dense `AttnImplType` with `--attn-impl`; FlashAttention 4 is the default | +| Video VAE | Supported | DiffVAE is the default; ConvVAE is selectable with `--video-vae conv` | +| CPU offload | Supported | `cpu` streams transformer blocks and releases modules between phases; `none` retains modules on the GPU | +| LoRA | Unsupported | The example does not expose a LoRA loader | +| Quantization | Unsupported | The example loads BF16 checkpoints | +| Feature cache | Unsupported | The example does not configure feature caching | +| Server API | Partial | Legacy `get_pipeline()` and `run_with_file()` entry points exist, but no explicit pipeline contract is declared | + +## Requirements + +- GPU: one, two, or four NVIDIA H100s; other GPU targets are not validated by this example +- Software: the standard TeleFuser environment, `transformers==5.14.1`, and an `ffmpeg` executable on `PATH` +- DiffVAE: a matching NATTEN/libnatten build is required for the formal 1536x1024, 121-frame workload +- I2V input: a PIL-readable still image such as PNG or JPEG + +Install TeleFuser by following the [development setup](../../CONTRIBUTING.md#development-setup). For the formal +DiffVAE path, select the command matching the installed PyTorch and CUDA versions from the +[NATTEN installation guide](https://natten.org/install/), then verify that its CUDA kernel library is available: + +```bash +python -c "import natten; print(natten.HAS_LIBNATTEN)" +``` + +The command must print `True`. Without NATTEN, the DiffVAE decoder uses the Triton/eager compatibility fallback, +which is not the formal performance or accuracy baseline. + +## Model Directory + +`--model-root` must point to the root of this exact split-checkpoint layout: + +```text +/path/to/LTX-2.5/ +|-- diffusion_models/ +| \-- ltx-2.5-22b-distilled-transformer-bf16.safetensors +|-- text_encoders/ +| \-- gemma4-12b-with-proj-ltx-2.5-bf16.safetensors +|-- vae/ +| |-- ltx-2.5-video-vae-bf16.safetensors +| |-- ltx-2.5-video-vae-conv-bf16.safetensors +| \-- ltx-2.5-audio-vae-bf16.safetensors +|-- latent_upscale_models/ +| \-- ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors +\-- model_patches/ + \-- ltx-2.5-duration-head-bf16.safetensors +``` + +The current checkpoint resolver validates all seven files at pipeline construction time, including both video VAE +checkpoints regardless of the `--video-vae` selection. + +Validate the split model pack without loading checkpoint tensors: + +```bash +python tools/validation/inspect_ltx25_checkpoints.py \ + --model-root /path/to/LTX-2.5 \ + --output work_dirs/ltx25-checkpoints.json +``` + +## Quick Start + +Run the default T2V workload from the repository root: + +```bash +python examples/ltx25_distilled/ltx25_distilled_t2v_h100.py \ + --model-root /path/to/LTX-2.5 \ + --prompt "A cinematic camera orbit around the subject." \ + --output-path work_dirs/ltx25-t2v.mp4 +``` + +The command writes a 1536x1024, 121-frame, 24 FPS video with synchronized audio to +`work_dirs/ltx25-t2v.mp4`. + +Run the denoising stage with four-way Ulysses sequence parallelism and FSDP2 shards: + +```bash +python examples/ltx25_distilled/ltx25_distilled_t2v_h100.py \ + --gpu-num 4 \ + --attn-impl FLASH_ATTN_4 \ + --model-root /path/to/LTX-2.5 \ + --prompt "A cinematic camera orbit around the subject." \ + --output-path work_dirs/ltx25-t2v-sp4.mp4 +``` + +## Examples + +### `ltx25_distilled_t2v_h100.py` + +This is the standalone text-to-video entry point. + +```bash +python examples/ltx25_distilled/ltx25_distilled_t2v_h100.py \ + --model-root /path/to/LTX-2.5 \ + --prompt "Ocean waves roll beneath a cloudy sky as distant thunder echoes." \ + --output-path work_dirs/ltx25-t2v.mp4 +``` + +### `ltx25_distilled_i2v_h100.py` + +This is the standalone image-to-video entry point. It uses the repository's frozen reference image by default; +`--image-path` can override it when reproducing another 896x512, 121-frame workload: + +```bash +python examples/ltx25_distilled/ltx25_distilled_i2v_h100.py \ + --model-root /path/to/LTX-2.5 \ + --image-path examples/data/ltx25/official_guitar_man.png \ + --image-frame-index 0 \ + --image-strength 1.0 \ + --prompt "A man with short gray hair plays a red electric guitar." \ + --output-path work_dirs/ltx25-i2v.mp4 +``` + +Key options: + +| Option | Default | Description | +| --- | --- | --- | +| `--prompt` | Script-specific | Text prompt used for video and audio generation | +| `--model-root` | Deployment-specific | Root of the required split model pack; pass it explicitly | +| `--output-path` | Required | Destination MP4; parent directories are created automatically | +| `--image-path` | `examples/data/ltx25/official_guitar_man.png` | Still image supplied to `ltx25_distilled_i2v_h100.py` | +| `--image-frame-index` | `0` (I2V) | Non-negative output-frame index for the image condition | +| `--image-strength` | `1.0` (I2V) | Image-conditioning strength in the inclusive range `[0, 1]` | +| `--height` | T2V: `1024`; I2V: `512` | Output height; must be a positive multiple of 64 | +| `--width` | T2V: `1536`; I2V: `896` | Output width; must be a positive multiple of 64 | +| `--num-frames` | `121` | Output frame count; must satisfy `num_frames = 8k + 1` | +| `--frame-rate` | `24.0` | Positive output frame rate in FPS | +| `--seed` | `42` | Random seed | +| `--video-vae` | `diff` | Video decoder: `diff` or `conv` | +| `--offload` | `cpu` | Model residency policy: `cpu` or `none` | +| `--gpu-num` | `1` | GPU count: `1`, `2`, or `4`; multi-GPU runs use Ulysses and FSDP2 | +| `--attn-impl` | `FLASH_ATTN_4` | Dense attention implementation from `AttnImplType` | + +Key behavior: + +- Both examples load checkpoint components into a CPU `ModuleManager`, then initialize six independently managed + stages for text encoding, video conditioning, denoising, latent upsampling, video decoding, and audio decoding. +- The distilled pipeline runs fixed two-stage sampling and jointly generates video and audio. +- Output audio is written as stereo PCM at 48 kHz and muxed into the MP4 as AAC. +- `LTX25DistilledOutput` also exposes the final video and audio VAE latents. +- Both examples support one, two, or four GPUs; only the I2V entry point accepts a still-image condition. + +## Regression + +`examples/run_examples.py` registers complete two-H100 T2V and I2V workloads. Both generate 121 frames, preserve the +audio stream, and use the runner's deterministic SDPA regression backend: + +```bash +python examples/run_examples.py --pipeline ltx25_distilled_t2v_2gpu --gpus 0,1 +python examples/run_examples.py --pipeline ltx25_distilled_i2v_2gpu --gpus 0,1 +``` + +Initialize or intentionally replace local baselines with `--update-baseline`. Normal regression runs require those +baselines and compare video PSNR/SSIM plus the audio stream contract and waveform similarity. + +## Configuration + +### Pipeline Composition + +`load_ltx25_distilled_modules()` constructs the split-checkpoint components on CPU and registers them with +`ModuleManager`. `LTX25DistilledPipeline.init()` composes the six flat stage modules from manager-owned components. +The compatibility constructor `LTX25DistilledPipeline.from_model_root()` follows the same loading and stage path. + +### Video VAE + +`--video-vae diff` selects DiffVAE and is the formal output path. It uses NATTEN when the compatible CUDA extension +is installed and otherwise falls back to the Triton/eager implementation. `--video-vae conv` selects ConvVAE as a +compatibility alternative. The selected VAE is used consistently for image conditioning, latent-statistics +normalization around spatial upsampling, and video decoding. Both variants return RGB chunks in `[F, H, W, C]` with +values in `[0, 1]`. Image conditions are applied to the clean latent state before initial noising. + +### Model Residency + +`--offload cpu` is the default. It streams transformer blocks between CPU and GPU, releases other modules at phase +boundaries, and lowers peak GPU residency at the cost of transfers. Transformer weights remain in pinned CPU memory +and stream through reusable GPU buffers; the other stage-owned modules move through GPU memory sequentially. Use +`--offload none` only when the GPU has enough memory to retain modules between phases. + +With `--gpu-num 2` or `--gpu-num 4`, the denoiser remains GPU-resident because FSDP2 and transformer CPU offload are +mutually exclusive. When `--offload cpu` is selected, the other five stages still follow their normal CPU-offload +lifecycle. + +### Sequence Parallelism + +`--gpu-num 2` and `--gpu-num 4` shard video and audio tokens across the denoising workers with Ulysses. Self-attention +and audio/video cross-attention exchange sequence and head partitions with all-to-all collectives; text context stays +replicated. The implementation pads non-divisible token counts and masks the padding before attention, then gathers +and crops model outputs before each sampler step. The 48 transformer blocks are also sharded with FSDP2. BF16 SP runs +are not bitwise-identical to single-GPU runs because the per-rank token dimensions can select different GEMM kernels. + +### Attention Backend + +`--attn-impl` selects a dense backend through `ModelRuntimeConfig.attention_config` and the public `telefuser.ops` +attention dispatcher. `FLASH_ATTN_4` preserves the validated single-H100 baseline. Other choices require their normal +runtime dependencies and use the dispatcher's documented fallback behavior where applicable. When SP must pad a +non-divisible sequence, only the affected masked attention calls use SDPA because the FlashAttention kernels do not +consume arbitrary additive masks. + +### Request Constraints + +- `height` and `width` must be positive multiples of 64. +- `num_frames` must satisfy `num_frames = 8k + 1`; examples include 1, 9, 17, and 121. +- `frame_rate` must be positive. +- `image_frame_index` must be non-negative and `image_strength` must be in `[0, 1]`. + +## Performance + +The formal quality and performance gates use BF16 on one H100 80 GB with PyTorch 2.11.0, CUDA 12.8, NATTEN 0.21.6, +the upstream eager DiffVAE tiling, and matching request/runtime provenance. The 1536x1024, 121-frame T2V comparison +recorded 61.90 dB PSNR and 0.999685 SSIM; the frozen 896x512, 121-frame I2V comparison recorded 62.95 dB PSNR and +0.999671 SSIM. + +The timings below are synchronized end-to-end p50 seconds from five cold and five warm samples: + +| Workload | Mode | Upstream cold / warm | TeleFuser cold / warm | +| --- | --- | ---: | ---: | +| T2V 1536x1024 / 121 | `offload=cpu` | 76.78 / 77.29 | 64.44 / 60.09 | +| I2V 896x512 / 121 | `offload=cpu` | 65.02 / 64.52 | 51.98 / 44.08 | +| T2V 1536x1024 / 121 | `offload=none` | 58.28 / 55.29 | 46.79 / 46.93 | +| I2V 896x512 / 121 | `offload=none` | 48.45 / 42.34 | 30.24 / 30.29 | + +The no-offload TeleFuser run reserved 79.14 GB at peak. Lower resolutions and shorter valid frame counts are useful +for diagnostics but do not replace these formal gates. + +## Troubleshooting + +### Missing Checkpoint + +The pipeline reports the first missing component and its resolved path. Compare `--model-root` with the complete +layout above; both video VAE files are currently required even when only one decoder is selected. + +### Output Has No Audio + +TeleFuser keeps the generated video if audio muxing fails. Confirm that `ffmpeg` is installed and available on +`PATH`: + +```bash +ffmpeg -version +``` + +### DiffVAE Uses the Compatibility Fallback + +Confirm that NATTEN is importable and includes libnatten for the active PyTorch/CUDA environment: + +```bash +python -c "import torch, natten; print(torch.__version__, torch.version.cuda, natten.HAS_LIBNATTEN)" +``` + +Use the [NATTEN installation matrix](https://natten.org/install/) to select a matching build when the final value is +`False`. + +## Notes + +- The formal workload is 1536x1024, 121 frames, 24 FPS, DiffVAE, and NATTEN on one H100. +- Lower resolutions and shorter valid frame counts are useful for smoke tests but are not the formal quality baseline. +- The frozen I2V input is `examples/data/ltx25/official_guitar_man.png`, used at frame index 0 and strength 1.0. +- CUDA audio-vocoder convolution can vary slightly across replays. The validation capture tools provide + `--deterministic-audio` for exact waveform comparisons without changing production inference behavior. diff --git a/examples/ltx25_distilled/ltx25_distilled_i2v_h100.py b/examples/ltx25_distilled/ltx25_distilled_i2v_h100.py new file mode 100644 index 0000000..bc1408d --- /dev/null +++ b/examples/ltx25_distilled/ltx25_distilled_i2v_h100.py @@ -0,0 +1,240 @@ +"""Run the faithful LTX-2.5 distilled image-to-video pipeline on H100 GPUs.""" + +from __future__ import annotations + +import tempfile +from pathlib import Path +from typing import Any, Literal, cast + +import click +import torch +from PIL import Image + +from telefuser.core.config import AttnImplType +from telefuser.core.module_manager import ModuleManager +from telefuser.pipelines.ltx25_distilled import ( + LTX25DistilledOutput, + LTX25DistilledPipeline, + LTX25ImageCondition, + build_ltx25_distilled_config, + load_ltx25_distilled_modules, +) +from telefuser.utils.audio import save_wav +from telefuser.utils.video import save_video + +DEFAULT_IMAGE_PATH = Path(__file__).resolve().parents[1] / "data" / "ltx25" / "official_guitar_man.png" + +PPL_CONFIG: dict[str, Any] = { + "name": "ltx25_distilled_i2v_h100", + "model_root": "/hhb-data/aigc/model_zoo/Lightricks/LTX-2.5/LTX-2.5", + "height": 512, + "width": 896, + "num_frames": 121, + "frame_rate": 24.0, + "seed": 42, + "prompt": "A man with short gray hair plays a red electric guitar.", + "input_image_path": str(DEFAULT_IMAGE_PATH), + "video_vae": "diff", + "attn_impl": "FLASH_ATTN_4", +} + +DENSE_ATTN_IMPLS = tuple( + implementation.name + for implementation in AttnImplType + if implementation not in {AttnImplType.RADIAL_ATTN, AttnImplType.LOCAL_SPARSE_ATTN, AttnImplType.SOL_ATTN} +) + + +def get_pipeline( + parallelism: int = 1, + model_root: str = PPL_CONFIG["model_root"], + video_vae: str = PPL_CONFIG["video_vae"], + offload: str = "cpu", + attn_impl: str | AttnImplType = PPL_CONFIG["attn_impl"], +) -> LTX25DistilledPipeline: + """Load the isolated LTX-2.5 distilled I2V pipeline on one or more H100s.""" + if parallelism not in (1, 2, 4): + raise ValueError(f"parallelism must be 1, 2, or 4, got {parallelism}") + if video_vae not in ("diff", "conv"): + raise ValueError(f"video_vae must be 'diff' or 'conv', got {video_vae!r}") + if offload not in ("none", "cpu"): + raise ValueError(f"offload must be 'none' or 'cpu', got {offload!r}") + selected_video_vae = cast(Literal["diff", "conv"], video_vae) + selected_offload = cast(Literal["none", "cpu"], offload) + if isinstance(attn_impl, str): + try: + selected_attn_impl = AttnImplType[attn_impl] + except KeyError as exc: + raise ValueError(f"Unknown attention implementation {attn_impl!r}; choose from {DENSE_ATTN_IMPLS}") from exc + else: + selected_attn_impl = attn_impl + if selected_attn_impl.name not in DENSE_ATTN_IMPLS: + raise ValueError(f"LTX-2.5 supports dense attention implementations only, got {selected_attn_impl.name}") + module_manager = ModuleManager(device="cpu", torch_dtype=torch.bfloat16) + load_ltx25_distilled_modules( + module_manager, + model_root, + video_vae=selected_video_vae, + torch_dtype=torch.bfloat16, + ) + pipeline = LTX25DistilledPipeline(device="cuda", torch_dtype=torch.bfloat16) + pipeline.init( + module_manager, + build_ltx25_distilled_config( + "cuda", + torch.bfloat16, + selected_video_vae, + selected_offload, + parallelism=parallelism, + attn_impl=selected_attn_impl, + ), + ) + return pipeline + + +def run( + pipeline: LTX25DistilledPipeline, + image: Image.Image, + prompt: str, + *, + seed: int = PPL_CONFIG["seed"], + height: int = PPL_CONFIG["height"], + width: int = PPL_CONFIG["width"], + num_frames: int = PPL_CONFIG["num_frames"], + frame_rate: float = PPL_CONFIG["frame_rate"], + image_frame_index: int = 0, + image_strength: float = 1.0, +) -> LTX25DistilledOutput: + """Generate video and synchronized audio from an image and text prompt.""" + images = ( + LTX25ImageCondition( + image.convert("RGB"), + frame_idx=image_frame_index, + strength=image_strength, + ), + ) + return pipeline( + prompt, + seed=seed, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + images=images, + ) + + +def run_with_file( + pipeline: LTX25DistilledPipeline, + prompt: str, + output_path: str, + input_image_path: str = PPL_CONFIG["input_image_path"], + seed: int = PPL_CONFIG["seed"], + height: int = PPL_CONFIG["height"], + width: int = PPL_CONFIG["width"], + num_frames: int = PPL_CONFIG["num_frames"], + frame_rate: float = PPL_CONFIG["frame_rate"], + image_frame_index: int = 0, + image_strength: float = 1.0, + first_image_path: str | None = None, + **_: object, +) -> dict[str, str]: + """Generate and save an MP4 with synchronized LTX-2.5 audio.""" + image = Image.open(first_image_path or input_image_path).convert("RGB") + result = run( + pipeline, + image, + prompt, + seed=seed, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + image_frame_index=image_frame_index, + image_strength=image_strength, + ) + frames = torch.cat(result.video_chunks).mul(255).round().clamp(0, 255).to(torch.uint8).cpu().numpy() + destination = Path(output_path) + destination.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as stream: + audio_path = Path(stream.name) + try: + save_wav(result.audio, 48000, str(audio_path)) + save_video(list(frames), str(destination), fps=result.frame_rate, quality=6, audio_path=str(audio_path)) + finally: + audio_path.unlink(missing_ok=True) + return {"output_path": str(destination)} + + +@click.command() +@click.option("--prompt", default=PPL_CONFIG["prompt"], show_default=True) +@click.option("--model-root", default=PPL_CONFIG["model_root"], show_default=True) +@click.option("--output-path", type=click.Path(path_type=Path), required=True) +@click.option( + "--image-path", + "input_image_path", + type=click.Path(exists=True, dir_okay=False, path_type=Path), + default=PPL_CONFIG["input_image_path"], + show_default=True, +) +@click.option("--image-frame-index", default=0, show_default=True) +@click.option("--image-strength", default=1.0, show_default=True) +@click.option("--height", default=PPL_CONFIG["height"], show_default=True) +@click.option("--width", default=PPL_CONFIG["width"], show_default=True) +@click.option("--num-frames", default=PPL_CONFIG["num_frames"], show_default=True) +@click.option("--frame-rate", default=PPL_CONFIG["frame_rate"], show_default=True) +@click.option("--seed", default=PPL_CONFIG["seed"], show_default=True) +@click.option("--video-vae", type=click.Choice(["diff", "conv"]), default=PPL_CONFIG["video_vae"], show_default=True) +@click.option("--offload", type=click.Choice(["none", "cpu"]), default="cpu", show_default=True) +@click.option("--gpu-num", type=click.Choice(["1", "2", "4"]), default="1", show_default=True) +@click.option( + "--attn-impl", + type=click.Choice(DENSE_ATTN_IMPLS), + default=PPL_CONFIG["attn_impl"], + show_default=True, +) +def main( + prompt: str, + model_root: str, + output_path: Path, + input_image_path: Path, + image_frame_index: int, + image_strength: float, + height: int, + width: int, + num_frames: int, + frame_rate: float, + seed: int, + video_vae: str, + offload: str, + gpu_num: str, + attn_impl: str, +) -> None: + """Generate an LTX-2.5 video and synchronized audio from an image.""" + pipeline = get_pipeline( + parallelism=int(gpu_num), + model_root=model_root, + video_vae=video_vae, + offload=offload, + attn_impl=attn_impl, + ) + try: + run_with_file( + pipeline, + input_image_path=str(input_image_path), + prompt=prompt, + output_path=str(output_path), + seed=seed, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + image_frame_index=image_frame_index, + image_strength=image_strength, + ) + finally: + pipeline.close() + + +if __name__ == "__main__": + main() diff --git a/examples/ltx25_distilled/ltx25_distilled_t2v_h100.py b/examples/ltx25_distilled/ltx25_distilled_t2v_h100.py new file mode 100644 index 0000000..1bb789e --- /dev/null +++ b/examples/ltx25_distilled/ltx25_distilled_t2v_h100.py @@ -0,0 +1,201 @@ +"""Run the faithful LTX-2.5 distilled text-to-video pipeline on H100 GPUs.""" + +from __future__ import annotations + +import tempfile +from pathlib import Path +from typing import Any, Literal, cast + +import click +import torch + +from telefuser.core.config import AttnImplType +from telefuser.core.module_manager import ModuleManager +from telefuser.pipelines.ltx25_distilled import ( + LTX25DistilledOutput, + LTX25DistilledPipeline, + build_ltx25_distilled_config, + load_ltx25_distilled_modules, +) +from telefuser.utils.audio import save_wav +from telefuser.utils.video import save_video + +PPL_CONFIG: dict[str, Any] = { + "name": "ltx25_distilled_t2v_h100", + "model_root": "/hhb-data/aigc/model_zoo/Lightricks/LTX-2.5/LTX-2.5", + "height": 1024, + "width": 1536, + "num_frames": 121, + "frame_rate": 24.0, + "seed": 42, + "prompt": "A cinematic camera orbit around the subject.", + "video_vae": "diff", + "attn_impl": "FLASH_ATTN_4", +} + +DENSE_ATTN_IMPLS = tuple( + implementation.name + for implementation in AttnImplType + if implementation not in {AttnImplType.RADIAL_ATTN, AttnImplType.LOCAL_SPARSE_ATTN, AttnImplType.SOL_ATTN} +) + + +def get_pipeline( + parallelism: int = 1, + model_root: str = PPL_CONFIG["model_root"], + video_vae: str = PPL_CONFIG["video_vae"], + offload: str = "cpu", + attn_impl: str | AttnImplType = PPL_CONFIG["attn_impl"], +) -> LTX25DistilledPipeline: + """Load the isolated LTX-2.5 distilled T2V pipeline on one or more H100s.""" + if parallelism not in (1, 2, 4): + raise ValueError(f"parallelism must be 1, 2, or 4, got {parallelism}") + if video_vae not in ("diff", "conv"): + raise ValueError(f"video_vae must be 'diff' or 'conv', got {video_vae!r}") + if offload not in ("none", "cpu"): + raise ValueError(f"offload must be 'none' or 'cpu', got {offload!r}") + selected_video_vae = cast(Literal["diff", "conv"], video_vae) + selected_offload = cast(Literal["none", "cpu"], offload) + if isinstance(attn_impl, str): + try: + selected_attn_impl = AttnImplType[attn_impl] + except KeyError as exc: + raise ValueError(f"Unknown attention implementation {attn_impl!r}; choose from {DENSE_ATTN_IMPLS}") from exc + else: + selected_attn_impl = attn_impl + if selected_attn_impl.name not in DENSE_ATTN_IMPLS: + raise ValueError(f"LTX-2.5 supports dense attention implementations only, got {selected_attn_impl.name}") + module_manager = ModuleManager(device="cpu", torch_dtype=torch.bfloat16) + load_ltx25_distilled_modules( + module_manager, + model_root, + video_vae=selected_video_vae, + torch_dtype=torch.bfloat16, + ) + pipeline = LTX25DistilledPipeline(device="cuda", torch_dtype=torch.bfloat16) + pipeline.init( + module_manager, + build_ltx25_distilled_config( + "cuda", + torch.bfloat16, + selected_video_vae, + selected_offload, + parallelism=parallelism, + attn_impl=selected_attn_impl, + ), + ) + return pipeline + + +def run( + pipeline: LTX25DistilledPipeline, + prompt: str, + *, + seed: int = PPL_CONFIG["seed"], + height: int = PPL_CONFIG["height"], + width: int = PPL_CONFIG["width"], + num_frames: int = PPL_CONFIG["num_frames"], + frame_rate: float = PPL_CONFIG["frame_rate"], +) -> LTX25DistilledOutput: + """Generate video and synchronized audio from a text prompt.""" + return pipeline( + prompt, + seed=seed, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + ) + + +def run_with_file( + pipeline: LTX25DistilledPipeline, + prompt: str, + output_path: str, + seed: int = PPL_CONFIG["seed"], + height: int = PPL_CONFIG["height"], + width: int = PPL_CONFIG["width"], + num_frames: int = PPL_CONFIG["num_frames"], + frame_rate: float = PPL_CONFIG["frame_rate"], + **_: object, +) -> dict[str, str]: + """Generate and save an MP4 with synchronized LTX-2.5 audio.""" + result = run( + pipeline, + prompt, + seed=seed, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + ) + frames = torch.cat(result.video_chunks).mul(255).round().clamp(0, 255).to(torch.uint8).cpu().numpy() + destination = Path(output_path) + destination.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as stream: + audio_path = Path(stream.name) + try: + save_wav(result.audio, 48000, str(audio_path)) + save_video(list(frames), str(destination), fps=result.frame_rate, quality=6, audio_path=str(audio_path)) + finally: + audio_path.unlink(missing_ok=True) + return {"output_path": str(destination)} + + +@click.command() +@click.option("--prompt", default=PPL_CONFIG["prompt"], show_default=True) +@click.option("--model-root", default=PPL_CONFIG["model_root"], show_default=True) +@click.option("--output-path", type=click.Path(path_type=Path), required=True) +@click.option("--height", default=PPL_CONFIG["height"], show_default=True) +@click.option("--width", default=PPL_CONFIG["width"], show_default=True) +@click.option("--num-frames", default=PPL_CONFIG["num_frames"], show_default=True) +@click.option("--frame-rate", default=PPL_CONFIG["frame_rate"], show_default=True) +@click.option("--seed", default=PPL_CONFIG["seed"], show_default=True) +@click.option("--video-vae", type=click.Choice(["diff", "conv"]), default=PPL_CONFIG["video_vae"], show_default=True) +@click.option("--offload", type=click.Choice(["none", "cpu"]), default="cpu", show_default=True) +@click.option("--gpu-num", type=click.Choice(["1", "2", "4"]), default="1", show_default=True) +@click.option( + "--attn-impl", + type=click.Choice(DENSE_ATTN_IMPLS), + default=PPL_CONFIG["attn_impl"], + show_default=True, +) +def main( + prompt: str, + model_root: str, + output_path: Path, + height: int, + width: int, + num_frames: int, + frame_rate: float, + seed: int, + video_vae: str, + offload: str, + gpu_num: str, + attn_impl: str, +) -> None: + """Generate an LTX-2.5 video and synchronized audio from text.""" + pipeline = get_pipeline( + parallelism=int(gpu_num), + model_root=model_root, + video_vae=video_vae, + offload=offload, + attn_impl=attn_impl, + ) + try: + run_with_file( + pipeline, + prompt, + str(output_path), + seed=seed, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + ) + finally: + pipeline.close() + + +if __name__ == "__main__": + main() diff --git a/examples/minimax_h3/minimax_h3_turbo_lora_h100.py b/examples/minimax_h3/minimax_h3_turbo_lora_h100.py index 1e0a74c..13647af 100644 --- a/examples/minimax_h3/minimax_h3_turbo_lora_h100.py +++ b/examples/minimax_h3/minimax_h3_turbo_lora_h100.py @@ -24,6 +24,7 @@ "lora_strength": 1.0, "num_inference_steps": 9, "prompt": "Steam rises from the ramen while the family talks in the background.", + "input_image_path": str(MINIMAX_H3_DEFAULT_FL2VA_IMAGE), "target_video_length": 8, "seed": 0, "flow_shift": 6.0, @@ -89,13 +90,13 @@ def run( seed: int = PPL_CONFIG["seed"], output_path: str = "minimax_h3_turbo_lora.mp4", target_video_length: float = PPL_CONFIG["target_video_length"], - image_path: str = str(MINIMAX_H3_DEFAULT_FL2VA_IMAGE), + input_image_path: str = PPL_CONFIG["input_image_path"], ) -> MiniMaxH3Generation: """Generate an image-conditioned Turbo H3 clip.""" result = pipeline( task="fl2va", prompt=prompt, - conditions=[{"type": "image", "role": "keyframe", "uri": image_path, "frame_index": 0}], + conditions=[{"type": "image", "role": "keyframe", "uri": input_image_path, "frame_index": 0}], target={ "short_edge": PPL_CONFIG["short_edge"], "aspect_ratio": PPL_CONFIG["aspect_ratio"], @@ -115,7 +116,8 @@ def run_with_file( seed: int = PPL_CONFIG["seed"], output_path: str = "minimax_h3_turbo_lora.mp4", target_video_length: float = PPL_CONFIG["target_video_length"], - image_path: str = str(MINIMAX_H3_DEFAULT_FL2VA_IMAGE), + input_image_path: str = PPL_CONFIG["input_image_path"], + first_image_path: str | None = None, **_: object, ) -> dict[str, str]: """Service-compatible wrapper that returns the generated output path.""" @@ -125,7 +127,7 @@ def run_with_file( seed=seed, output_path=output_path, target_video_length=target_video_length, - image_path=image_path, + input_image_path=first_image_path or input_image_path, ) return {"output_path": output_path} @@ -135,7 +137,7 @@ def main() -> None: parser.add_argument("--gpu-num", "--ulysses-degree", dest="gpu_num", type=int, choices=(1, 2, 4), default=1) parser.add_argument("--model-root", default=PPL_CONFIG["model_root"]) parser.add_argument("--lora-path", default=PPL_CONFIG["lora_path"]) - parser.add_argument("--image", default=str(MINIMAX_H3_DEFAULT_FL2VA_IMAGE)) + parser.add_argument("--image", dest="input_image_path", default=PPL_CONFIG["input_image_path"]) parser.add_argument("--prompt", default=PPL_CONFIG["prompt"]) parser.add_argument("--duration", type=float, default=PPL_CONFIG["target_video_length"]) parser.add_argument("--steps", type=int, default=PPL_CONFIG["num_inference_steps"]) @@ -168,7 +170,7 @@ def main() -> None: seed=args.seed, output_path=args.output, target_video_length=args.duration, - image_path=args.image, + input_image_path=args.input_image_path, ) finally: pipeline.stop() diff --git a/examples/run_examples.py b/examples/run_examples.py index 6548fd1..e1f33bc 100644 --- a/examples/run_examples.py +++ b/examples/run_examples.py @@ -934,6 +934,8 @@ def get_param(param_name: str, default: str | None = None) -> str | None: from PIL import Image kwargs["image"] = Image.open(config["input_image_path"]).convert("RGB") + if "image_path" in params and config.get("input_image_path"): + kwargs["image_path"] = config["input_image_path"] if "audio_path" in params: kwargs["audio_path"] = config.get("input_audio_path") diff --git a/pyproject.toml b/pyproject.toml index 45973c9..6c2c1ca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,7 +64,7 @@ dependencies = [ "torchcodec>=0.7.0", "torchvision>=0.21.0", "tqdm>=4.66.0", - "transformers==4.57.3", + "transformers==5.14.1", "triton", "uvicorn>=0.30,<1.0", "torchao", diff --git a/telefuser/models/ltx25/__init__.py b/telefuser/models/ltx25/__init__.py new file mode 100644 index 0000000..800fb56 --- /dev/null +++ b/telefuser/models/ltx25/__init__.py @@ -0,0 +1,77 @@ +"""Isolated LTX-2.5 model support. + +The package deliberately does not import the LTX-2.3 model implementations. +""" + +from .audio import ( + LTX25AudioVAEDecoder, + LTX25AudioVocoder, + load_ltx25_audio_decoder_and_vocoder, + ltx25_audio_checkpoint_key_coverage, +) +from .checkpoint import ( + LTX25_COMPONENT_NAMES, + LTX25CheckpointMetadata, + LTX25ModelPaths, + inspect_checkpoint, + inspect_model_pack, +) +from .conv_video_vae import LTX25ConvVideoVAE, ltx25_conv_video_vae_checkpoint_key_coverage +from .diff_vae import DiffusionVideoDecoder +from .diff_vae.diffusion_video_decoder import ltx25_diffusion_vae_checkpoint_key_coverage +from .duration import LTX25DurationHead, ltx25_duration_checkpoint_key_coverage, seconds_to_num_frames +from .embeddings import LTX25EmbeddingsProcessor, LTX25EmbeddingsProcessorOutput +from .gemma4 import LTX25Gemma4TextEncoder, LTX25GemmaAssets, LTX25GemmaTokenizer +from .sampler import ( + LTX25_STAGE1_DISTILLED_SIGMAS, + LTX25_STAGE2_DISTILLED_SIGMAS, + LTX25EulerAncestralStep, + uses_ancestral_stage1_sampler, +) +from .spatial_upsampler import ( + LTX25PerChannelStatistics, + LTX25SpatialUpsampler, + LTX25SpatialUpsamplerConfig, + load_video_latent_statistics, + upsample_video_latent, +) +from .transformer import LTX25AVTransformer, build_ltx25_av_model, ltx25_transformer_key_to_model_key +from .video_encoder import LTX25VideoEncoder, ltx25_video_encoder_checkpoint_key_coverage + +__all__ = [ + "LTX25_COMPONENT_NAMES", + "LTX25ModelPaths", + "LTX25CheckpointMetadata", + "LTX25ConvVideoVAE", + "LTX25DurationHead", + "DiffusionVideoDecoder", + "LTX25AudioVAEDecoder", + "LTX25AudioVocoder", + "LTX25EmbeddingsProcessor", + "LTX25EmbeddingsProcessorOutput", + "LTX25_STAGE1_DISTILLED_SIGMAS", + "LTX25_STAGE2_DISTILLED_SIGMAS", + "LTX25EulerAncestralStep", + "LTX25Gemma4TextEncoder", + "LTX25GemmaAssets", + "LTX25GemmaTokenizer", + "LTX25PerChannelStatistics", + "LTX25SpatialUpsampler", + "LTX25SpatialUpsamplerConfig", + "LTX25AVTransformer", + "LTX25VideoEncoder", + "inspect_checkpoint", + "inspect_model_pack", + "build_ltx25_av_model", + "ltx25_transformer_key_to_model_key", + "ltx25_video_encoder_checkpoint_key_coverage", + "ltx25_audio_checkpoint_key_coverage", + "ltx25_conv_video_vae_checkpoint_key_coverage", + "ltx25_duration_checkpoint_key_coverage", + "ltx25_diffusion_vae_checkpoint_key_coverage", + "load_video_latent_statistics", + "load_ltx25_audio_decoder_and_vocoder", + "upsample_video_latent", + "uses_ancestral_stage1_sampler", + "seconds_to_num_frames", +] diff --git a/telefuser/models/ltx25/audio.py b/telefuser/models/ltx25/audio.py new file mode 100644 index 0000000..7a826f9 --- /dev/null +++ b/telefuser/models/ltx25/audio.py @@ -0,0 +1,1300 @@ +"""Isolated LTX-2.5 audio VAE decoder, vocoder, and split-checkpoint loader.""" + +from __future__ import annotations + +import contextlib +import math +from collections.abc import Iterator +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from typing import Any + +import einops +import torch +import torch.nn.functional as F +from safetensors import safe_open +from torch import nn + +from .checkpoint import inspect_checkpoint + + +@contextlib.contextmanager +def _module_in_fp32(module: nn.Module, *, enabled: bool) -> Iterator[None]: + """Temporarily materialize a module in fp32 when autocast cannot do so.""" + if not enabled: + yield + return + module_dtype = next(module.parameters()).dtype + module.float() + try: + yield + finally: + module.to(module_dtype) + + +# ----------------------------------------------------------------------------- +# Common normalization +# ----------------------------------------------------------------------------- + + +class NormType(Enum): + """Normalization layer types: GROUP (GroupNorm) or PIXEL (per-location RMS norm).""" + + GROUP = "group" + PIXEL = "pixel" + + +class PixelNorm(nn.Module): + """Per-location RMS normalization.""" + + def __init__(self, dim: int = 1, eps: float = 1e-8) -> None: + super().__init__() + self.dim = dim + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True) + rms = torch.sqrt(mean_sq + self.eps) + return x / rms + + +def build_normalization_layer( + in_channels: int, + *, + num_groups: int = 32, + normtype: NormType = NormType.GROUP, +) -> nn.Module: + if normtype == NormType.GROUP: + return nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True) + if normtype == NormType.PIXEL: + return PixelNorm(dim=1, eps=1e-6) + raise ValueError(f"Invalid normalization type: {normtype}") + + +# ----------------------------------------------------------------------------- +# Audio VAE blocks (decoder only) +# ----------------------------------------------------------------------------- + + +class CausalityAxis(Enum): + """Enum for specifying the causality axis in causal convolutions.""" + + NONE = None + WIDTH = "width" + HEIGHT = "height" + WIDTH_COMPATIBILITY = "width-compatibility" + + +class CausalConv2d(nn.Module): + """A causal 2D convolution wrapper used by LTX audio VAE.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int], + stride: int = 1, + dilation: int | tuple[int, int] = 1, + groups: int = 1, + bias: bool = True, + causality_axis: CausalityAxis = CausalityAxis.HEIGHT, + ) -> None: + super().__init__() + self.causality_axis = causality_axis + kernel_size = nn.modules.utils._pair(kernel_size) + dilation = nn.modules.utils._pair(dilation) + + pad_h = (kernel_size[0] - 1) * dilation[0] + pad_w = (kernel_size[1] - 1) * dilation[1] + + match self.causality_axis: + case CausalityAxis.NONE: + self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2) + case CausalityAxis.WIDTH | CausalityAxis.WIDTH_COMPATIBILITY: + self.padding = (pad_w, 0, pad_h // 2, pad_h - pad_h // 2) + case CausalityAxis.HEIGHT: + self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h, 0) + case _: + raise ValueError(f"Invalid causality_axis: {causality_axis}") + + self.conv = nn.Conv2d( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=0, + dilation=dilation, + groups=groups, + bias=bias, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = F.pad(x, self.padding) + return self.conv(x) + + +def make_conv2d( + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int], + stride: int = 1, + padding: tuple[int, int, int, int] | None = None, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + causality_axis: CausalityAxis | None = None, +) -> nn.Module: + if causality_axis is not None: + return CausalConv2d(in_channels, out_channels, kernel_size, stride, dilation, groups, bias, causality_axis) + if padding is None: + padding = kernel_size // 2 if isinstance(kernel_size, int) else tuple(k // 2 for k in kernel_size) + return nn.Conv2d( + in_channels, + out_channels, + kernel_size, + stride, + padding, + dilation, + groups, + bias, + ) + + +class AttentionType(Enum): + """Enum for specifying the attention mechanism type.""" + + VANILLA = "vanilla" + LINEAR = "linear" + NONE = "none" + + +class AttnBlock(nn.Module): + def __init__( + self, + in_channels: int, + norm_type: NormType = NormType.GROUP, + ) -> None: + super().__init__() + self.in_channels = in_channels + + self.norm = build_normalization_layer(in_channels, normtype=norm_type) + self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h_ = self.norm(x) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + b, c, h, w = q.shape + q = q.reshape(b, c, h * w).contiguous().permute(0, 2, 1).contiguous() + k = k.reshape(b, c, h * w).contiguous() + w_ = torch.bmm(q, k).contiguous() + w_ = w_ * (int(c) ** (-0.5)) + w_ = torch.nn.functional.softmax(w_, dim=2) + + v = v.reshape(b, c, h * w).contiguous() + w_ = w_.permute(0, 2, 1).contiguous() + h_ = torch.bmm(v, w_).contiguous().reshape(b, c, h, w).contiguous() + h_ = self.proj_out(h_) + return x + h_ + + +def make_attn( + in_channels: int, + attn_type: AttentionType = AttentionType.VANILLA, + norm_type: NormType = NormType.GROUP, +) -> nn.Module: + match attn_type: + case AttentionType.VANILLA: + return AttnBlock(in_channels, norm_type=norm_type) + case AttentionType.NONE: + return nn.Identity() + case AttentionType.LINEAR: + raise NotImplementedError(f"Attention type {attn_type.value} is not supported yet.") + case _: + raise ValueError(f"Unknown attention type: {attn_type}") + + +LRELU_SLOPE = 0.1 + + +class ResBlock1(nn.Module): + """1D ResBlock used by the vocoder.""" + + def __init__(self, channels: int, kernel_size: int = 3, dilation: tuple[int, int, int] = (1, 3, 5)) -> None: + super().__init__() + self.convs1 = nn.ModuleList( + [ + nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], padding="same"), + nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], padding="same"), + nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2], padding="same"), + ] + ) + self.convs2 = nn.ModuleList( + [ + nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding="same"), + nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding="same"), + nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding="same"), + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for conv1, conv2 in zip(self.convs1, self.convs2, strict=True): + xt = torch.nn.functional.leaky_relu(x, LRELU_SLOPE) + xt = conv1(xt) + xt = torch.nn.functional.leaky_relu(xt, LRELU_SLOPE) + xt = conv2(xt) + x = xt + x + return x + + +class ResnetBlock(nn.Module): + """2D ResNet block used by the audio decoder.""" + + def __init__( # noqa: PLR0913 + self, + *, + in_channels: int, + out_channels: int | None = None, + conv_shortcut: bool = False, + dropout: float = 0.0, + temb_channels: int = 512, + norm_type: NormType = NormType.GROUP, + causality_axis: CausalityAxis = CausalityAxis.HEIGHT, + ) -> None: + super().__init__() + self.causality_axis = causality_axis + + if self.causality_axis != CausalityAxis.NONE and norm_type == NormType.GROUP: + raise ValueError("Causal ResnetBlock with GroupNorm is not supported.") + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.use_conv_shortcut = conv_shortcut + + self.norm1 = build_normalization_layer(in_channels, normtype=norm_type) + self.non_linearity = nn.SiLU() + self.conv1 = make_conv2d(in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis) + if temb_channels > 0: + self.temb_proj = nn.Linear(temb_channels, out_channels) + self.norm2 = build_normalization_layer(out_channels, normtype=norm_type) + self.dropout = nn.Dropout(dropout) + self.conv2 = make_conv2d(out_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis) + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + self.conv_shortcut = make_conv2d( + in_channels, + out_channels, + kernel_size=3, + stride=1, + causality_axis=causality_axis, + ) + else: + self.nin_shortcut = make_conv2d( + in_channels, + out_channels, + kernel_size=1, + stride=1, + causality_axis=causality_axis, + ) + + def forward(self, x: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor: + h = self.norm1(x) + h = self.non_linearity(h) + h = self.conv1(h) + if temb is not None: + h = h + self.temb_proj(self.non_linearity(temb))[:, :, None, None] + h = self.norm2(h) + h = self.non_linearity(h) + h = self.dropout(h) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + x = self.conv_shortcut(x) if self.use_conv_shortcut else self.nin_shortcut(x) + return x + h + + +class Upsample(nn.Module): + def __init__( + self, + in_channels: int, + with_conv: bool, + causality_axis: CausalityAxis = CausalityAxis.HEIGHT, + ) -> None: + super().__init__() + self.with_conv = with_conv + self.causality_axis = causality_axis + if self.with_conv: + self.conv = make_conv2d(in_channels, in_channels, kernel_size=3, stride=1, causality_axis=causality_axis) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") + if self.with_conv: + x = self.conv(x) + match self.causality_axis: + case CausalityAxis.NONE: + pass + case CausalityAxis.HEIGHT: + x = x[:, :, 1:, :] + case CausalityAxis.WIDTH: + x = x[:, :, :, 1:] + case CausalityAxis.WIDTH_COMPATIBILITY: + pass + case _: + raise ValueError(f"Invalid causality_axis: {self.causality_axis}") + return x + + +def build_upsampling_path( # noqa: PLR0913 + *, + ch: int, + ch_mult: tuple[int, ...], + num_resolutions: int, + num_res_blocks: int, + resolution: int, + temb_channels: int, + dropout: float, + norm_type: NormType, + causality_axis: CausalityAxis, + attn_type: AttentionType, + attn_resolutions: set[int], + resamp_with_conv: bool, + initial_block_channels: int, +) -> tuple[nn.ModuleList, int]: + up_modules = nn.ModuleList() + block_in = initial_block_channels + curr_res = resolution // (2 ** (num_resolutions - 1)) + + for level in reversed(range(num_resolutions)): + stage = nn.Module() + stage.block = nn.ModuleList() + stage.attn = nn.ModuleList() + block_out = ch * ch_mult[level] + + for _ in range(num_res_blocks + 1): + stage.block.append( + ResnetBlock( + in_channels=block_in, + out_channels=block_out, + temb_channels=temb_channels, + dropout=dropout, + norm_type=norm_type, + causality_axis=causality_axis, + ) + ) + block_in = block_out + if curr_res in attn_resolutions: + stage.attn.append(make_attn(block_in, attn_type=attn_type, norm_type=norm_type)) + + if level != 0: + stage.upsample = Upsample(block_in, resamp_with_conv, causality_axis=causality_axis) + curr_res *= 2 + + up_modules.insert(0, stage) + + return up_modules, block_in + + +class AudioPatchifier: + """Patchifier tailored for spectrogram/audio latents (used for normalization only).""" + + def __init__(self, patch_size: int) -> None: + self.patch_size = (patch_size, 1, 1) + + def patchify(self, latents: torch.Tensor) -> torch.Tensor: + return einops.rearrange(latents, "b c (f p) m -> b f (c p m)", p=self.patch_size[0]) + + def unpatchify(self, latents: torch.Tensor, output_shape: "AudioLatentShape") -> torch.Tensor: + return einops.rearrange( + latents, + "b f (c p m) -> b c (f p) m", + c=output_shape.channels, + p=self.patch_size[0], + m=output_shape.mel_bins, + ) + + +@dataclass(frozen=True) +class AudioLatentShape: + batch: int + channels: int + frames: int + mel_bins: int + + +class PerChannelStatistics(nn.Module): + """Per-channel statistics for denormalizing the latent representation.""" + + def __init__(self, latent_channels: int = 128) -> None: + super().__init__() + self.register_buffer("std-of-means", torch.empty(latent_channels)) + self.register_buffer("mean-of-means", torch.empty(latent_channels)) + + def un_normalize(self, x: torch.Tensor) -> torch.Tensor: + return (x * self.get_buffer("std-of-means").to(x)) + self.get_buffer("mean-of-means").to(x) + + +def build_mid_block( # noqa: PLR0913 + channels: int, + temb_channels: int, + dropout: float, + norm_type: NormType, + causality_axis: CausalityAxis, + attn_type: AttentionType, + add_attention: bool, +) -> nn.Module: + mid = nn.Module() + mid.block_1 = ResnetBlock( + in_channels=channels, + out_channels=channels, + temb_channels=temb_channels, + dropout=dropout, + norm_type=norm_type, + causality_axis=causality_axis, + ) + mid.attn_1 = make_attn(channels, attn_type=attn_type, norm_type=norm_type) if add_attention else nn.Identity() + mid.block_2 = ResnetBlock( + in_channels=channels, + out_channels=channels, + temb_channels=temb_channels, + dropout=dropout, + norm_type=norm_type, + causality_axis=causality_axis, + ) + return mid + + +def run_mid_block(mid: nn.Module, features: torch.Tensor) -> torch.Tensor: + features = mid.block_1(features, temb=None) + features = mid.attn_1(features) + return mid.block_2(features, temb=None) + + +class AudioDecoder(nn.Module): + """Decoder that reconstructs audio spectrograms from latent features.""" + + def __init__( # noqa: PLR0913 + self, + *, + ch: int, + out_ch: int, + ch_mult: tuple[int, ...] = (1, 2, 4, 8), + num_res_blocks: int = 2, + attn_resolutions: set[int] | None = None, + dropout: float = 0.0, + in_channels: int = 2, + resolution: int = 256, + z_channels: int = 8, + double_z: bool = True, + attn_type: AttentionType = AttentionType.VANILLA, + mid_block_add_attention: bool = True, + norm_type: NormType = NormType.PIXEL, + causality_axis: CausalityAxis = CausalityAxis.HEIGHT, + sample_rate: int = 16000, + mel_hop_length: int = 160, + is_causal: bool = True, + mel_bins: int | None = None, + **_ignore_kwargs: Any, + ) -> None: + super().__init__() + _ = (in_channels, double_z) + resamp_with_conv = True + if attn_resolutions is None: + attn_resolutions = {8, 16, 32} + + self.per_channel_statistics = PerChannelStatistics(latent_channels=ch) + self.sample_rate = sample_rate + self.mel_hop_length = mel_hop_length + self.is_causal = is_causal + self.mel_bins = mel_bins + self.patchifier = AudioPatchifier( + patch_size=1, + ) + + self.ch = ch + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.out_ch = out_ch + self.give_pre_end = False + self.tanh_out = False + self.norm_type = norm_type + self.z_channels = z_channels + self.channel_multipliers = ch_mult + self.attn_resolutions = attn_resolutions + self.causality_axis = causality_axis + self.attn_type = attn_type + + base_block_channels = ch * self.channel_multipliers[-1] + self.conv_in = make_conv2d( + z_channels, + base_block_channels, + kernel_size=3, + stride=1, + causality_axis=self.causality_axis, + ) + self.non_linearity = nn.SiLU() + self.mid = build_mid_block( + channels=base_block_channels, + temb_channels=self.temb_ch, + dropout=dropout, + norm_type=self.norm_type, + causality_axis=self.causality_axis, + attn_type=self.attn_type, + add_attention=mid_block_add_attention, + ) + self.up, final_block_channels = build_upsampling_path( + ch=ch, + ch_mult=ch_mult, + num_resolutions=self.num_resolutions, + num_res_blocks=num_res_blocks, + resolution=resolution, + temb_channels=self.temb_ch, + dropout=dropout, + norm_type=self.norm_type, + causality_axis=self.causality_axis, + attn_type=self.attn_type, + attn_resolutions=self.attn_resolutions, + resamp_with_conv=resamp_with_conv, + initial_block_channels=base_block_channels, + ) + self.norm_out = build_normalization_layer(final_block_channels, normtype=self.norm_type) + self.conv_out = make_conv2d( + final_block_channels, + out_ch, + kernel_size=3, + stride=1, + causality_axis=self.causality_axis, + ) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + sample, target_shape = self._denormalize_latents(sample) + h = self.conv_in(sample) + h = run_mid_block(self.mid, h) + h = self._run_upsampling_path(h) + h = self._finalize_output(h) + return self._adjust_output_shape(h, target_shape) + + def _denormalize_latents(self, sample: torch.Tensor) -> tuple[torch.Tensor, AudioLatentShape]: + latent_shape = AudioLatentShape( + batch=sample.shape[0], + channels=sample.shape[1], + frames=sample.shape[2], + mel_bins=sample.shape[3], + ) + sample_patched = self.patchifier.patchify(sample) + sample_denormalized = self.per_channel_statistics.un_normalize(sample_patched) + sample = self.patchifier.unpatchify(sample_denormalized, latent_shape) + + target_frames = latent_shape.frames * 4 + if self.causality_axis != CausalityAxis.NONE: + target_frames = max(target_frames - (4 - 1), 1) + target_shape = AudioLatentShape( + batch=latent_shape.batch, + channels=self.out_ch, + frames=target_frames, + mel_bins=self.mel_bins if self.mel_bins is not None else latent_shape.mel_bins, + ) + return sample, target_shape + + def _run_upsampling_path(self, h: torch.Tensor) -> torch.Tensor: + for level in reversed(range(self.num_resolutions)): + stage = self.up[level] + for block_idx, block in enumerate(stage.block): + h = block(h, temb=None) + if stage.attn: + h = stage.attn[block_idx](h) + if level != 0 and hasattr(stage, "upsample"): + h = stage.upsample(h) + return h + + def _finalize_output(self, h: torch.Tensor) -> torch.Tensor: + if self.give_pre_end: + return h + h = self.norm_out(h) + h = self.non_linearity(h) + h = self.conv_out(h) + return torch.tanh(h) if self.tanh_out else h + + def _adjust_output_shape(self, decoded_output: torch.Tensor, target_shape: AudioLatentShape) -> torch.Tensor: + _, _, current_time, current_freq = decoded_output.shape + target_channels = target_shape.channels + target_time = target_shape.frames + target_freq = target_shape.mel_bins + + decoded_output = decoded_output[ + :, :target_channels, : min(current_time, target_time), : min(current_freq, target_freq) + ] + time_padding_needed = target_time - decoded_output.shape[2] + freq_padding_needed = target_freq - decoded_output.shape[3] + if time_padding_needed > 0 or freq_padding_needed > 0: + padding = (0, max(freq_padding_needed, 0), 0, max(time_padding_needed, 0)) + decoded_output = F.pad(decoded_output, padding) + + decoded_output = decoded_output[:, :target_channels, :target_time, :target_freq] + return decoded_output + + +# ----------------------------------------------------------------------------- +# Vocoder (copied from LTX-2) +# ----------------------------------------------------------------------------- + + +def get_padding(kernel_size: int, dilation: int = 1) -> int: + return int((kernel_size * dilation - dilation) / 2) + + +def _sinc(x: torch.Tensor) -> torch.Tensor: + return torch.where( + x == 0, + torch.tensor(1.0, device=x.device, dtype=x.dtype), + torch.sin(math.pi * x) / math.pi / x, + ) + + +def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor: + even = kernel_size % 2 == 0 + half_size = kernel_size // 2 + delta_f = 4 * half_width + amplitude = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 + if amplitude > 50.0: + beta = 0.1102 * (amplitude - 8.7) + elif amplitude >= 21.0: + beta = 0.5842 * (amplitude - 21) ** 0.4 + 0.07886 * (amplitude - 21.0) + else: + beta = 0.0 + window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) + time = torch.arange(-half_size, half_size) + 0.5 if even else torch.arange(kernel_size) - half_size + if cutoff == 0: + filter_ = torch.zeros_like(time) + else: + filter_ = 2 * cutoff * window * _sinc(2 * cutoff * time) + filter_ /= filter_.sum() + return filter_.view(1, 1, kernel_size) + + +class LowPassFilter1d(nn.Module): + def __init__( + self, + cutoff: float = 0.5, + half_width: float = 0.6, + stride: int = 1, + padding: bool = True, + padding_mode: str = "replicate", + kernel_size: int = 12, + ) -> None: + super().__init__() + if cutoff < -0.0: + raise ValueError("Minimum cutoff must be larger than zero.") + if cutoff > 0.5: + raise ValueError("A cutoff above 0.5 does not make sense.") + self.kernel_size = kernel_size + self.even = kernel_size % 2 == 0 + self.pad_left = kernel_size // 2 - int(self.even) + self.pad_right = kernel_size // 2 + self.stride = stride + self.padding = padding + self.padding_mode = padding_mode + self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + _, n_channels, _ = x.shape + if self.padding: + x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode) + filt = self.filter.to(dtype=x.dtype, device=x.device).expand(n_channels, -1, -1) + return F.conv1d(x, filt, stride=self.stride, groups=n_channels) + + +class UpSample1d(nn.Module): + def __init__( + self, + ratio: int = 2, + kernel_size: int | None = None, + persistent: bool = True, + window_type: str = "kaiser", + ) -> None: + super().__init__() + self.ratio = ratio + self.stride = ratio + + if window_type == "hann": + rolloff = 0.99 + lowpass_filter_width = 6 + width = math.ceil(lowpass_filter_width / rolloff) + self.kernel_size = 2 * width * ratio + 1 + self.pad = width + self.pad_left = 2 * width * ratio + self.pad_right = self.kernel_size - ratio + time_axis = (torch.arange(self.kernel_size) / ratio - width) * rolloff + time_clamped = time_axis.clamp(-lowpass_filter_width, lowpass_filter_width) + window = torch.cos(time_clamped * math.pi / lowpass_filter_width / 2) ** 2 + sinc_filter = (torch.sinc(time_axis) * window * rolloff / ratio).view(1, 1, -1) + else: + self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + self.pad = self.kernel_size // ratio - 1 + self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 + self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 + sinc_filter = kaiser_sinc_filter1d( + cutoff=0.5 / ratio, + half_width=0.6 / ratio, + kernel_size=self.kernel_size, + ) + + self.register_buffer("filter", sinc_filter, persistent=persistent) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + _, n_channels, _ = x.shape + x = F.pad(x, (self.pad, self.pad), mode="replicate") + filt = self.filter.to(dtype=x.dtype, device=x.device).expand(n_channels, -1, -1) + x = self.ratio * F.conv_transpose1d(x, filt, stride=self.stride, groups=n_channels) + return x[..., self.pad_left : -self.pad_right] + + +class DownSample1d(nn.Module): + def __init__(self, ratio: int = 2, kernel_size: int | None = None) -> None: + super().__init__() + self.ratio = ratio + self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + self.lowpass = LowPassFilter1d( + cutoff=0.5 / ratio, + half_width=0.6 / ratio, + stride=ratio, + kernel_size=self.kernel_size, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.lowpass(x) + + +class Snake(nn.Module): + def __init__( + self, + in_features: int, + alpha: float = 1.0, + alpha_trainable: bool = True, + alpha_logscale: bool = True, + ) -> None: + super().__init__() + self.alpha_logscale = alpha_logscale + self.alpha = nn.Parameter(torch.zeros(in_features) if alpha_logscale else torch.ones(in_features) * alpha) + self.alpha.requires_grad = alpha_trainable + self.eps = 1e-9 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) + if self.alpha_logscale: + alpha = torch.exp(alpha) + return x + (1.0 / (alpha + self.eps)) * torch.sin(x * alpha).pow(2) + + +class SnakeBeta(nn.Module): + def __init__( + self, + in_features: int, + alpha: float = 1.0, + alpha_trainable: bool = True, + alpha_logscale: bool = True, + ) -> None: + super().__init__() + self.alpha_logscale = alpha_logscale + self.alpha = nn.Parameter(torch.zeros(in_features) if alpha_logscale else torch.ones(in_features) * alpha) + self.alpha.requires_grad = alpha_trainable + self.beta = nn.Parameter(torch.zeros(in_features) if alpha_logscale else torch.ones(in_features) * alpha) + self.beta.requires_grad = alpha_trainable + self.eps = 1e-9 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) + beta = self.beta.unsqueeze(0).unsqueeze(-1) + if self.alpha_logscale: + alpha = torch.exp(alpha) + beta = torch.exp(beta) + return x + (1.0 / (beta + self.eps)) * torch.sin(x * alpha).pow(2) + + +class Activation1d(nn.Module): + def __init__( + self, + activation: nn.Module, + up_ratio: int = 2, + down_ratio: int = 2, + up_kernel_size: int = 12, + down_kernel_size: int = 12, + ) -> None: + super().__init__() + self.act = activation + self.upsample = UpSample1d(up_ratio, up_kernel_size) + self.downsample = DownSample1d(down_ratio, down_kernel_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.upsample(x) + x = self.act(x) + return self.downsample(x) + + +class AMPBlock1(nn.Module): + def __init__( + self, + channels: int, + kernel_size: int = 3, + dilation: tuple[int, int, int] = (1, 3, 5), + activation: str = "snake", + ) -> None: + super().__init__() + act_cls = SnakeBeta if activation == "snakebeta" else Snake + self.convs1 = nn.ModuleList( + [ + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + padding=get_padding(kernel_size, dilation[0]), + ), + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + padding=get_padding(kernel_size, dilation[1]), + ), + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[2], + padding=get_padding(kernel_size, dilation[2]), + ), + ] + ) + + self.convs2 = nn.ModuleList( + [ + nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)), + nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)), + nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)), + ] + ) + + self.acts1 = nn.ModuleList([Activation1d(act_cls(channels)) for _ in range(len(self.convs1))]) + self.acts2 = nn.ModuleList([Activation1d(act_cls(channels)) for _ in range(len(self.convs2))]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for c1, c2, a1, a2 in zip(self.convs1, self.convs2, self.acts1, self.acts2, strict=True): + xt = a1(x) + xt = c1(xt) + xt = a2(xt) + xt = c2(xt) + x = x + xt + return x + + +class Vocoder(nn.Module): + def __init__( # noqa: PLR0913 + self, + resblock_kernel_sizes: list[int] | None = None, + upsample_rates: list[int] | None = None, + upsample_kernel_sizes: list[int] | None = None, + resblock_dilation_sizes: list[list[int]] | None = None, + upsample_initial_channel: int = 1024, + resblock: str = "1", + output_sampling_rate: int = 24000, + activation: str = "snake", + use_tanh_at_final: bool = True, + apply_final_activation: bool = True, + use_bias_at_final: bool = True, + ) -> None: + super().__init__() + if resblock_kernel_sizes is None: + resblock_kernel_sizes = [3, 7, 11] + if upsample_rates is None: + upsample_rates = [6, 5, 2, 2, 2] + if upsample_kernel_sizes is None: + upsample_kernel_sizes = [16, 15, 8, 4, 4] + if resblock_dilation_sizes is None: + resblock_dilation_sizes = [[1, 3, 5], [1, 3, 5], [1, 3, 5]] + + self.output_sampling_rate = output_sampling_rate + self.num_kernels = len(resblock_kernel_sizes) + self.num_upsamples = len(upsample_rates) + self.use_tanh_at_final = use_tanh_at_final + self.apply_final_activation = apply_final_activation + self.is_amp = resblock == "AMP1" + + self.conv_pre = nn.Conv1d( + in_channels=128, + out_channels=upsample_initial_channel, + kernel_size=7, + stride=1, + padding=3, + ) + + self.ups = nn.ModuleList( + nn.ConvTranspose1d( + upsample_initial_channel // (2**i), + upsample_initial_channel // (2 ** (i + 1)), + kernel_size, + stride, + padding=(kernel_size - stride) // 2, + ) + for i, (stride, kernel_size) in enumerate(zip(upsample_rates, upsample_kernel_sizes, strict=True)) + ) + + final_channels = upsample_initial_channel // (2 ** len(upsample_rates)) + self.resblocks = nn.ModuleList() + for i in range(len(upsample_rates)): + ch = upsample_initial_channel // (2 ** (i + 1)) + for kernel_size, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes, strict=True): + if self.is_amp: + self.resblocks.append(AMPBlock1(ch, kernel_size, tuple(dilations), activation=activation)) + else: + self.resblocks.append(ResBlock1(ch, kernel_size, tuple(dilations))) + + if self.is_amp: + self.act_post: nn.Module = Activation1d(SnakeBeta(final_channels)) + else: + self.act_post = nn.LeakyReLU() + + self.conv_post = nn.Conv1d( + in_channels=final_channels, + out_channels=2, + kernel_size=7, + stride=1, + padding=3, + bias=use_bias_at_final, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x.transpose(2, 3) + if x.dim() == 4: + if x.shape[1] != 2: + raise ValueError("Input must have 2 channels for stereo") + x = einops.rearrange(x, "b s c t -> b (s c) t") + + x = self.conv_pre(x) + for i in range(self.num_upsamples): + if not self.is_amp: + x = F.leaky_relu(x, LRELU_SLOPE) + x = self.ups[i](x) + start = i * self.num_kernels + end = start + self.num_kernels + block_outputs = torch.stack([self.resblocks[idx](x) for idx in range(start, end)], dim=0) + x = block_outputs.mean(dim=0) + + x = self.act_post(x) + x = self.conv_post(x) + if self.apply_final_activation: + x = torch.tanh(x) if self.use_tanh_at_final else torch.clamp(x, -1, 1) + return x + + +class _STFTFn(nn.Module): + def __init__(self, filter_length: int, hop_length: int, win_length: int) -> None: + super().__init__() + self.hop_length = hop_length + self.win_length = win_length + n_freqs = filter_length // 2 + 1 + self.register_buffer("forward_basis", torch.zeros(n_freqs * 2, 1, filter_length)) + self.register_buffer("inverse_basis", torch.zeros(n_freqs * 2, 1, filter_length)) + + def forward(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + if y.dim() == 2: + y = y.unsqueeze(1) + left_pad = max(0, self.win_length - self.hop_length) + y = F.pad(y, (left_pad, 0)) + spec = F.conv1d(y, self.forward_basis, stride=self.hop_length, padding=0) + n_freqs = spec.shape[1] // 2 + real, imag = spec[:, :n_freqs], spec[:, n_freqs:] + magnitude = torch.sqrt(real**2 + imag**2) + phase = torch.atan2(imag.float(), real.float()).to(real.dtype) + return magnitude, phase + + +class MelSTFT(nn.Module): + """Causal log-mel spectrogram module whose buffers are loaded from the checkpoint.""" + + def __init__(self, filter_length: int, hop_length: int, win_length: int, n_mel_channels: int) -> None: + super().__init__() + self.stft_fn = _STFTFn(filter_length, hop_length, win_length) + n_freqs = filter_length // 2 + 1 + self.register_buffer("mel_basis", torch.zeros(n_mel_channels, n_freqs)) + + def mel_spectrogram(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + magnitude, phase = self.stft_fn(y) + energy = torch.norm(magnitude, dim=1) + mel = torch.matmul(self.mel_basis.to(magnitude.dtype), magnitude) + log_mel = torch.log(torch.clamp(mel, min=1e-5)) + return log_mel, magnitude, phase, energy + + +class VocoderWithBWE(nn.Module): + def __init__( + self, + vocoder: Vocoder, + bwe_generator: Vocoder, + mel_stft: MelSTFT, + input_sampling_rate: int, + output_sampling_rate: int, + hop_length: int, + ) -> None: + super().__init__() + self.vocoder = vocoder + self.bwe_generator = bwe_generator + self.mel_stft = mel_stft + self.input_sampling_rate = input_sampling_rate + self.output_sampling_rate = output_sampling_rate + self.hop_length = hop_length + + # Resampler filter is not stored in checkpoint (persistent=False). Build on CPU to materialize buffers. + with torch.device("cpu"): + self.resampler = UpSample1d( + ratio=output_sampling_rate // input_sampling_rate, + persistent=False, + window_type="hann", + ) + + @property + def conv_pre(self) -> nn.Conv1d: + return self.vocoder.conv_pre + + @property + def conv_post(self) -> nn.Conv1d: + return self.vocoder.conv_post + + def _compute_mel(self, audio: torch.Tensor) -> torch.Tensor: + batch, n_channels, _ = audio.shape + flat = audio.reshape(batch * n_channels, -1) + mel, _, _, _ = self.mel_stft.mel_spectrogram(flat) + return mel.reshape(batch, n_channels, mel.shape[1], mel.shape[2]) + + def forward(self, mel_spec: torch.Tensor) -> torch.Tensor: + """Decode in fp32 to match the upstream BWE accumulation contract.""" + input_dtype = mel_spec.dtype + device_type = mel_spec.device.type + module_dtype = next(self.parameters()).dtype + fp32_ctx = ( + _module_in_fp32(self, enabled=module_dtype != torch.float32) + if device_type == "mps" + else torch.autocast(device_type=device_type, dtype=torch.float32) + ) + with fp32_ctx: + out = self.vocoder(mel_spec.float()) + _, _, length_low_rate = out.shape + output_length = length_low_rate * self.output_sampling_rate // self.input_sampling_rate + remainder = length_low_rate % self.hop_length + if remainder != 0: + out = F.pad(out, (0, self.hop_length - remainder)) + mel = self._compute_mel(out) + residual = self.bwe_generator(mel.transpose(2, 3)) + skip = self.resampler(out) + if residual.shape != skip.shape: + raise ValueError(f"residual {residual.shape} != skip {skip.shape}") + return torch.clamp(residual + skip, -1, 1)[..., :output_length].to(input_dtype) + + +# ----------------------------------------------------------------------------- +# Configurators (subset of ltx_core model_configurator) +# ----------------------------------------------------------------------------- + + +class AudioDecoderConfigurator: + @classmethod + def from_config(cls, config: dict[str, Any]) -> AudioDecoder: + audio_vae_cfg = config.get("audio_vae", {}) + model_cfg = audio_vae_cfg.get("model", {}) + model_params = model_cfg.get("params", {}) + ddconfig = model_params.get("ddconfig", {}) + preprocessing_cfg = audio_vae_cfg.get("preprocessing", {}) + stft_cfg = preprocessing_cfg.get("stft", {}) + mel_cfg = preprocessing_cfg.get("mel", {}) + variables_cfg = audio_vae_cfg.get("variables", {}) + + sample_rate = int(model_params.get("sampling_rate", 16000)) + mel_hop_length = int(stft_cfg.get("hop_length", 160)) + is_causal = bool(stft_cfg.get("causal", True)) + mel_bins = ddconfig.get("mel_bins") or mel_cfg.get("n_mel_channels") or variables_cfg.get("mel_bins") + mel_bins = int(mel_bins) if mel_bins is not None else None + + return AudioDecoder( + ch=int(ddconfig.get("ch", 128)), + out_ch=int(ddconfig.get("out_ch", 2)), + ch_mult=tuple(ddconfig.get("ch_mult", (1, 2, 4))), + num_res_blocks=int(ddconfig.get("num_res_blocks", 2)), + attn_resolutions=set(ddconfig.get("attn_resolutions", {8, 16, 32})), + resolution=int(ddconfig.get("resolution", 256)), + z_channels=int(ddconfig.get("z_channels", 8)), + norm_type=NormType(ddconfig.get("norm_type", "pixel")), + causality_axis=CausalityAxis(ddconfig.get("causality_axis", "height")), + dropout=float(ddconfig.get("dropout", 0.0)), + mid_block_add_attention=bool(ddconfig.get("mid_block_add_attention", True)), + sample_rate=sample_rate, + mel_hop_length=mel_hop_length, + is_causal=is_causal, + mel_bins=mel_bins, + ) + + +class VocoderConfigurator: + @classmethod + def from_config(cls, config: dict[str, Any]) -> Vocoder | VocoderWithBWE: + cfg = config.get("vocoder", {}) + if "bwe" not in cfg: + return Vocoder( + resblock_kernel_sizes=cfg.get("resblock_kernel_sizes", [3, 7, 11]), + upsample_rates=cfg.get("upsample_rates", [6, 5, 2, 2, 2]), + upsample_kernel_sizes=cfg.get("upsample_kernel_sizes", [16, 15, 8, 4, 4]), + resblock_dilation_sizes=cfg.get("resblock_dilation_sizes", [[1, 3, 5], [1, 3, 5], [1, 3, 5]]), + upsample_initial_channel=int(cfg.get("upsample_initial_channel", 1024)), + resblock=str(cfg.get("resblock", "1")), + output_sampling_rate=int(cfg.get("output_sampling_rate", 24000)), + activation=str(cfg.get("activation", "snake")), + use_tanh_at_final=bool(cfg.get("use_tanh_at_final", True)), + apply_final_activation=True, + use_bias_at_final=bool(cfg.get("use_bias_at_final", True)), + ) + + vocoder_cfg = cfg.get("vocoder", {}) + bwe_cfg = cfg["bwe"] + vocoder = Vocoder( + resblock_kernel_sizes=vocoder_cfg.get("resblock_kernel_sizes", [3, 7, 11]), + upsample_rates=vocoder_cfg.get("upsample_rates", [6, 5, 2, 2, 2]), + upsample_kernel_sizes=vocoder_cfg.get("upsample_kernel_sizes", [16, 15, 8, 4, 4]), + resblock_dilation_sizes=vocoder_cfg.get("resblock_dilation_sizes", [[1, 3, 5], [1, 3, 5], [1, 3, 5]]), + upsample_initial_channel=int(vocoder_cfg.get("upsample_initial_channel", 1024)), + resblock=str(vocoder_cfg.get("resblock", "AMP1")), + output_sampling_rate=int(bwe_cfg.get("input_sampling_rate", 24000)), + activation=str(vocoder_cfg.get("activation", "snakebeta")), + use_tanh_at_final=bool(vocoder_cfg.get("use_tanh_at_final", True)), + apply_final_activation=True, + use_bias_at_final=bool(vocoder_cfg.get("use_bias_at_final", True)), + ) + bwe_generator = Vocoder( + resblock_kernel_sizes=bwe_cfg.get("resblock_kernel_sizes", [3, 7, 11]), + upsample_rates=bwe_cfg.get("upsample_rates", [6, 5, 2, 2, 2]), + upsample_kernel_sizes=bwe_cfg.get("upsample_kernel_sizes", [16, 15, 8, 4, 4]), + resblock_dilation_sizes=bwe_cfg.get("resblock_dilation_sizes", [[1, 3, 5], [1, 3, 5], [1, 3, 5]]), + upsample_initial_channel=int(bwe_cfg.get("upsample_initial_channel", 1024)), + resblock=str(bwe_cfg.get("resblock", "AMP1")), + output_sampling_rate=int(bwe_cfg.get("output_sampling_rate", 48000)), + activation=str(bwe_cfg.get("activation", "snakebeta")), + use_tanh_at_final=bool(bwe_cfg.get("use_tanh_at_final", True)), + apply_final_activation=False, + use_bias_at_final=bool(bwe_cfg.get("use_bias_at_final", True)), + ) + mel_stft = MelSTFT( + filter_length=int(bwe_cfg.get("n_fft", 1024)), + hop_length=int(bwe_cfg.get("hop_length", 256)), + win_length=int(bwe_cfg.get("n_fft", 1024)), + n_mel_channels=int(bwe_cfg.get("num_mels", 64)), + ) + return VocoderWithBWE( + vocoder=vocoder, + bwe_generator=bwe_generator, + mel_stft=mel_stft, + input_sampling_rate=int(bwe_cfg.get("input_sampling_rate", 24000)), + output_sampling_rate=int(bwe_cfg.get("output_sampling_rate", 48000)), + hop_length=int(bwe_cfg.get("hop_length", 256)), + ) + + +LTX25AudioVAEDecoder = AudioDecoder +LTX25AudioVocoder = Vocoder | VocoderWithBWE + + +def _audio_decoder_key_to_model_key(key: str) -> str | None: + if key.startswith("audio_vae.decoder."): + return key.removeprefix("audio_vae.decoder.") + if key.startswith("audio_vae.per_channel_statistics."): + return "per_channel_statistics." + key.removeprefix("audio_vae.per_channel_statistics.") + return None + + +def _vocoder_key_to_model_key(key: str) -> str | None: + if not key.startswith("vocoder."): + return None + return key.removeprefix("vocoder.") + + +def ltx25_audio_checkpoint_key_coverage( + checkpoint_path: str | Path, + decoder_keys: set[str], + vocoder_keys: set[str], +) -> tuple[set[str], set[str], set[str], set[str]]: + """Return unexpected and missing decoder/vocoder keys without loading payloads.""" + decoder_mapped: set[str] = set() + vocoder_mapped: set[str] = set() + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as checkpoint: + for key in checkpoint.keys(): + if (target := _audio_decoder_key_to_model_key(key)) is not None: + decoder_mapped.add(target) + if (target := _vocoder_key_to_model_key(key)) is not None: + vocoder_mapped.add(target) + return ( + decoder_mapped - decoder_keys, + decoder_keys - decoder_mapped, + vocoder_mapped - vocoder_keys, + vocoder_keys - vocoder_mapped, + ) + + +def load_ltx25_audio_decoder_and_vocoder( + checkpoint_path: str | Path, + *, + device: torch.device | str = "cpu", + torch_dtype: torch.dtype = torch.bfloat16, +) -> tuple[LTX25AudioVAEDecoder, LTX25AudioVocoder]: + """Construct and strictly load the LTX-2.5 audio decoder and BWE vocoder.""" + checkpoint = inspect_checkpoint(checkpoint_path) + with torch.device("meta"): + decoder = AudioDecoderConfigurator.from_config(checkpoint.config) + vocoder = VocoderConfigurator.from_config(checkpoint.config) + decoder_unexpected, decoder_missing, vocoder_unexpected, vocoder_missing = ltx25_audio_checkpoint_key_coverage( + checkpoint.path, set(decoder.state_dict()), set(vocoder.state_dict()) + ) + if decoder_unexpected or decoder_missing or vocoder_unexpected or vocoder_missing: + raise ValueError( + "LTX-2.5 audio checkpoint coverage mismatch: " + f"decoder_unexpected={sorted(decoder_unexpected)[:5]}, " + f"decoder_missing={sorted(decoder_missing)[:5]}, " + f"vocoder_unexpected={sorted(vocoder_unexpected)[:5]}, " + f"vocoder_missing={sorted(vocoder_missing)[:5]}" + ) + + decoder_state_dict: dict[str, torch.Tensor] = {} + vocoder_state_dict: dict[str, torch.Tensor] = {} + with safe_open(str(checkpoint.path), framework="pt", device="cpu") as source: + for key in source.keys(): + if (target := _audio_decoder_key_to_model_key(key)) is not None: + decoder_state_dict[target] = source.get_tensor(key) + if (target := _vocoder_key_to_model_key(key)) is not None: + vocoder_state_dict[target] = source.get_tensor(key) + decoder_missing_keys, decoder_unexpected_keys = decoder.load_state_dict( + decoder_state_dict, strict=True, assign=True + ) + vocoder_missing_keys, vocoder_unexpected_keys = vocoder.load_state_dict( + vocoder_state_dict, strict=True, assign=True + ) + if decoder_missing_keys or decoder_unexpected_keys or vocoder_missing_keys or vocoder_unexpected_keys: + raise ValueError( + "LTX-2.5 audio weight load mismatch: " + f"decoder_missing={decoder_missing_keys[:5]}, decoder_unexpected={decoder_unexpected_keys[:5]}, " + f"vocoder_missing={vocoder_missing_keys[:5]}, vocoder_unexpected={vocoder_unexpected_keys[:5]}" + ) + return decoder.to(device=device, dtype=torch_dtype).eval(), vocoder.to(device=device, dtype=torch_dtype).eval() + + +__all__ = [ + "LTX25AudioVAEDecoder", + "LTX25AudioVocoder", + "load_ltx25_audio_decoder_and_vocoder", + "ltx25_audio_checkpoint_key_coverage", +] diff --git a/telefuser/models/ltx25/checkpoint.py b/telefuser/models/ltx25/checkpoint.py new file mode 100644 index 0000000..c4f7d22 --- /dev/null +++ b/telefuser/models/ltx25/checkpoint.py @@ -0,0 +1,175 @@ +"""LTX-2.5 split-checkpoint discovery and metadata validation. + +This module owns the LTX-2.5 file layout. It is intentionally independent from +the legacy monolithic LTX checkpoint loader and does not instantiate model +classes; architecture construction must consume the returned metadata rather +than copy constants from LTX-2.3. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +from safetensors import safe_open + +LTX25_COMPONENT_NAMES = ( + "ltx25_transformer", + "ltx25_gemma4", + "ltx25_embeddings_processor", + "ltx25_video_encoder", + "ltx25_video_decoder", + "ltx25_audio_decoder", + "ltx25_vocoder", + "ltx25_spatial_upsampler", + "ltx25_duration_head", +) + +_DEFAULT_PATHS = { + "transformer_path": "diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", + "text_encoder_path": "text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", + "video_vae_path": "vae/ltx-2.5-video-vae-bf16.safetensors", + "conv_video_vae_path": "vae/ltx-2.5-video-vae-conv-bf16.safetensors", + "audio_vae_path": "vae/ltx-2.5-audio-vae-bf16.safetensors", + "spatial_upsampler_path": "latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", + "duration_head_path": "model_patches/ltx-2.5-duration-head-bf16.safetensors", +} + + +def parse_model_version(version: str | None) -> tuple[int, ...]: + """Parse an LTX metadata version into comparable numeric components.""" + if not version: + return () + values: list[int] = [] + for part in version.replace("-", ".").split("."): + if not part.isdigit(): + break + values.append(int(part)) + return tuple(values) + + +def _require_file(path: Path, label: str) -> Path: + if not path.is_file(): + raise FileNotFoundError(f"LTX-2.5 {label} checkpoint does not exist: {path}") + return path + + +@dataclass(frozen=True, slots=True) +class LTX25ModelPaths: + """Resolved LTX-2.5 split-checkpoint paths.""" + + transformer_path: Path + text_encoder_path: Path + video_vae_path: Path + conv_video_vae_path: Path + audio_vae_path: Path + spatial_upsampler_path: Path + duration_head_path: Path + + @classmethod + def from_model_root(cls, model_root: str | Path) -> "LTX25ModelPaths": + """Resolve the official BF16 LTX-2.5 distilled model-pack layout.""" + root = Path(model_root).expanduser().resolve() + paths = { + name: _require_file(root / relative_path, name.removesuffix("_path")) + for name, relative_path in _DEFAULT_PATHS.items() + } + return cls(**paths) + + def as_dict(self) -> dict[str, str]: + """Return absolute paths in a JSON-compatible representation.""" + return {name: str(path) for name, path in asdict(self).items()} + + +@dataclass(frozen=True, slots=True) +class LTX25CheckpointMetadata: + """Provenance and unmodified safetensors metadata for one component.""" + + path: Path + size_bytes: int + sha256: str | None + tensor_count: int + metadata: dict[str, Any] + config: dict[str, Any] + model_version: tuple[int, ...] + + def as_dict(self) -> dict[str, Any]: + return { + "path": str(self.path), + "size_bytes": self.size_bytes, + "sha256": self.sha256, + "tensor_count": self.tensor_count, + "metadata": self.metadata, + "config": self.config, + "model_version": list(self.model_version), + } + + +def _decode_metadata_value(value: str) -> Any: + try: + return json.loads(value) + except json.JSONDecodeError: + return value + + +def sha256_file(path: Path) -> str: + """Return the SHA-256 of a checkpoint without loading tensor payloads.""" + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def inspect_checkpoint(path: str | Path, *, include_sha256: bool = False) -> LTX25CheckpointMetadata: + """Read metadata and key count without materializing checkpoint tensors.""" + resolved = _require_file(Path(path).expanduser().resolve(), "component") + with safe_open(resolved, framework="pt", device="cpu") as checkpoint: + raw_metadata = checkpoint.metadata() or {} + metadata = {key: _decode_metadata_value(value) for key, value in raw_metadata.items()} + tensor_count = len(checkpoint.keys()) + config = metadata.get("config", {}) + if not isinstance(config, dict): + raise ValueError(f"LTX-2.5 checkpoint config metadata must be an object: {resolved}") + metadata_version = metadata.get("model_version") + model_version = parse_model_version(metadata_version if isinstance(metadata_version, str) else None) + return LTX25CheckpointMetadata( + path=resolved, + size_bytes=resolved.stat().st_size, + sha256=sha256_file(resolved) if include_sha256 else None, + tensor_count=tensor_count, + metadata=metadata, + config=config, + model_version=model_version, + ) + + +def validate_gemma_source_checkpoint( + transformer_metadata: LTX25CheckpointMetadata, + gemma_config: dict[str, Any], +) -> None: + """Validate the Gemma4 version declared by an LTX-2.5 transformer checkpoint.""" + source = transformer_metadata.metadata.get("gemma_source_checkpoint") + if not isinstance(source, dict): + raise ValueError( + f"LTX-2.5 transformer {transformer_metadata.path} is missing gemma_source_checkpoint metadata." + ) + expected = source.get("gemma_version") + actual = gemma_config.get("gemma_version") + if expected != actual: + raise ValueError( + "Gemma version mismatch: transformer metadata expects " + f"gemma_version={expected!r}, but the Gemma config declares {actual!r}." + ) + + +def inspect_model_pack(model_root: str | Path, *, include_sha256: bool = False) -> dict[str, LTX25CheckpointMetadata]: + """Inspect every checkpoint required by the LTX-2.5 distilled reference path.""" + paths = LTX25ModelPaths.from_model_root(model_root) + return { + name.removesuffix("_path"): inspect_checkpoint(path, include_sha256=include_sha256) + for name, path in asdict(paths).items() + } diff --git a/telefuser/models/ltx25/conv_video_vae.py b/telefuser/models/ltx25/conv_video_vae.py new file mode 100644 index 0000000..3d07a96 --- /dev/null +++ b/telefuser/models/ltx25/conv_video_vae.py @@ -0,0 +1,2754 @@ +"""Isolated LTX-2.5 convolutional video VAE implementation.""" + +from __future__ import annotations + +from enum import Enum +from pathlib import Path +from typing import Iterable + +import torch +from safetensors import safe_open +from torch import nn + +from telefuser.core.base_model import BaseModel + +from .checkpoint import inspect_checkpoint + +DEFAULT_ENCODER_BLOCKS = ( + ("res_x", {"num_layers": 4}), + ("compress_space_res", {"multiplier": 2}), + ("res_x", {"num_layers": 6}), + ("compress_time_res", {"multiplier": 2}), + ("res_x", {"num_layers": 4}), + ("compress_all_res", {"multiplier": 2}), + ("res_x", {"num_layers": 2}), + ("compress_all_res", {"multiplier": 1}), + ("res_x", {"num_layers": 2}), +) + +DEFAULT_DECODER_BLOCKS = ( + ("res_x", {"num_layers": 4}), + ("compress_space", {"multiplier": 2}), + ("res_x", {"num_layers": 6}), + ("compress_time", {"multiplier": 2}), + ("res_x", {"num_layers": 4}), + ("compress_all", {"multiplier": 1}), + ("res_x", {"num_layers": 2}), + ("compress_all", {"multiplier": 2}), + ("res_x", {"num_layers": 2}), +) + + +class NormType(Enum): + """Normalization layer types: GROUP (GroupNorm) or PIXEL (per-location RMS norm).""" + + GROUP = "group" + PIXEL = "pixel" + + +class PixelNorm(nn.Module): + """ + Per-pixel (per-location) RMS normalization layer. + For each element along the chosen dimension, this layer normalizes the tensor + by the root-mean-square of its values across that dimension: + y = x / sqrt(mean(x^2, dim=dim, keepdim=True) + eps) + """ + + def __init__(self, dim: int = 1, eps: float = 1e-8) -> None: + """ + Args: + dim: Dimension along which to compute the RMS (typically channels). + eps: Small constant added for numerical stability. + """ + super().__init__() + self.dim = dim + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Apply RMS normalization along the configured dimension. + """ + # Compute mean of squared values along `dim`, keep dimensions for broadcasting. + mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True) + # Normalize by the root-mean-square (RMS). + rms = torch.sqrt(mean_sq + self.eps) + return x / rms + + +def build_normalization_layer( + in_channels: int, *, num_groups: int = 32, normtype: NormType = NormType.GROUP +) -> nn.Module: + """ + Create a normalization layer based on the normalization type. + Args: + in_channels: Number of input channels + num_groups: Number of groups for group normalization + normtype: Type of normalization: "group" or "pixel" + Returns: + A normalization layer + """ + if normtype == NormType.GROUP: + return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True) + if normtype == NormType.PIXEL: + return PixelNorm(dim=1, eps=1e-6) + raise ValueError(f"Invalid normalization type: {normtype}") + + +from enum import Enum + + +class NormLayerType(Enum): + GROUP_NORM = "group_norm" + PIXEL_NORM = "pixel_norm" + + +class LogVarianceType(Enum): + PER_CHANNEL = "per_channel" + UNIFORM = "uniform" + CONSTANT = "constant" + NONE = "none" + + +class PaddingModeType(Enum): + ZEROS = "zeros" + REFLECT = "reflect" + REPLICATE = "replicate" + CIRCULAR = "circular" + + +from typing import Tuple, Union + +from einops import rearrange +from torch import nn +from torch.nn import functional as F + + +def make_conv_nd( # noqa: PLR0913 + dims: Union[int, Tuple[int, int]], + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + padding: int = 0, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + causal: bool = False, + spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, + temporal_padding_mode: PaddingModeType = PaddingModeType.ZEROS, +) -> nn.Module: + if not (spatial_padding_mode == temporal_padding_mode or causal): + raise NotImplementedError("spatial and temporal padding modes must be equal") + if dims == 2: + return nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + bias=bias, + padding_mode=spatial_padding_mode.value, + ) + elif dims == 3: + if causal: + return CausalConv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + groups=groups, + bias=bias, + spatial_padding_mode=spatial_padding_mode, + ) + return nn.Conv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + bias=bias, + padding_mode=spatial_padding_mode.value, + ) + elif dims == (2, 1): + return DualConv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + bias=bias, + padding_mode=spatial_padding_mode.value, + ) + else: + raise ValueError(f"unsupported dimensions: {dims}") + + +def make_linear_nd( + dims: int, + in_channels: int, + out_channels: int, + bias: bool = True, +) -> nn.Module: + if dims == 2: + return nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias) + elif dims in (3, (2, 1)): + return nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias) + else: + raise ValueError(f"unsupported dimensions: {dims}") + + +class DualConv3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: Union[int, Tuple[int, int, int]] = 1, + padding: Union[int, Tuple[int, int, int]] = 0, + dilation: Union[int, Tuple[int, int, int]] = 1, + groups: int = 1, + bias: bool = True, + padding_mode: str = "zeros", + ) -> None: + super(DualConv3d, self).__init__() + + self.in_channels = in_channels + self.out_channels = out_channels + self.padding_mode = padding_mode + # Ensure kernel_size, stride, padding, and dilation are tuples of length 3 + if isinstance(kernel_size, int): + kernel_size = (kernel_size, kernel_size, kernel_size) + if kernel_size == (1, 1, 1): + raise ValueError("kernel_size must be greater than 1. Use make_linear_nd instead.") + if isinstance(stride, int): + stride = (stride, stride, stride) + if isinstance(padding, int): + padding = (padding, padding, padding) + if isinstance(dilation, int): + dilation = (dilation, dilation, dilation) + + # Set parameters for convolutions + self.groups = groups + self.bias = bias + + # Define the size of the channels after the first convolution + intermediate_channels = out_channels if in_channels < out_channels else in_channels + + # Define parameters for the first convolution + self.weight1 = nn.Parameter( + torch.Tensor( + intermediate_channels, + in_channels // groups, + 1, + kernel_size[1], + kernel_size[2], + ) + ) + self.stride1 = (1, stride[1], stride[2]) + self.padding1 = (0, padding[1], padding[2]) + self.dilation1 = (1, dilation[1], dilation[2]) + if bias: + self.bias1 = nn.Parameter(torch.Tensor(intermediate_channels)) + else: + self.register_parameter("bias1", None) + + # Define parameters for the second convolution + self.weight2 = nn.Parameter(torch.Tensor(out_channels, intermediate_channels // groups, kernel_size[0], 1, 1)) + self.stride2 = (stride[0], 1, 1) + self.padding2 = (padding[0], 0, 0) + self.dilation2 = (dilation[0], 1, 1) + if bias: + self.bias2 = nn.Parameter(torch.Tensor(out_channels)) + else: + self.register_parameter("bias2", None) + + # Initialize weights and biases + self.reset_parameters() + + def reset_parameters(self) -> None: + nn.init.kaiming_uniform_(self.weight1, a=torch.sqrt(5)) + nn.init.kaiming_uniform_(self.weight2, a=torch.sqrt(5)) + if self.bias: + fan_in1, _ = nn.init._calculate_fan_in_and_fan_out(self.weight1) + bound1 = 1 / torch.sqrt(fan_in1) + nn.init.uniform_(self.bias1, -bound1, bound1) + fan_in2, _ = nn.init._calculate_fan_in_and_fan_out(self.weight2) + bound2 = 1 / torch.sqrt(fan_in2) + nn.init.uniform_(self.bias2, -bound2, bound2) + + def forward( + self, + x: torch.Tensor, + use_conv3d: bool = False, + skip_time_conv: bool = False, + ) -> torch.Tensor: + if use_conv3d: + return self.forward_with_3d(x=x, skip_time_conv=skip_time_conv) + else: + return self.forward_with_2d(x=x, skip_time_conv=skip_time_conv) + + def forward_with_3d(self, x: torch.Tensor, skip_time_conv: bool = False) -> torch.Tensor: + # First convolution + x = F.conv3d( + x, + self.weight1, + self.bias1, + self.stride1, + self.padding1, + self.dilation1, + self.groups, + padding_mode=self.padding_mode, + ) + + if skip_time_conv: + return x + + # Second convolution + x = F.conv3d( + x, + self.weight2, + self.bias2, + self.stride2, + self.padding2, + self.dilation2, + self.groups, + padding_mode=self.padding_mode, + ) + + return x + + def forward_with_2d(self, x: torch.Tensor, skip_time_conv: bool = False) -> torch.Tensor: + b, _, _, h, w = x.shape + + # First 2D convolution + x = rearrange(x, "b c d h w -> (b d) c h w") + # Squeeze the depth dimension out of weight1 since it's 1 + weight1 = self.weight1.squeeze(2) + # Select stride, padding, and dilation for the 2D convolution + stride1 = (self.stride1[1], self.stride1[2]) + padding1 = (self.padding1[1], self.padding1[2]) + dilation1 = (self.dilation1[1], self.dilation1[2]) + x = F.conv2d( + x, + weight1, + self.bias1, + stride1, + padding1, + dilation1, + self.groups, + padding_mode=self.padding_mode, + ) + + _, _, h, w = x.shape + + if skip_time_conv: + x = rearrange(x, "(b d) c h w -> b c d h w", b=b) + return x + + # Second convolution which is essentially treated as a 1D convolution across the 'd' dimension + x = rearrange(x, "(b d) c h w -> (b h w) c d", b=b) + + # Reshape weight2 to match the expected dimensions for conv1d + weight2 = self.weight2.squeeze(-1).squeeze(-1) + # Use only the relevant dimension for stride, padding, and dilation for the 1D convolution + stride2 = self.stride2[0] + padding2 = self.padding2[0] + dilation2 = self.dilation2[0] + x = F.conv1d( + x, + weight2, + self.bias2, + stride2, + padding2, + dilation2, + self.groups, + padding_mode=self.padding_mode, + ) + x = rearrange(x, "(b h w) c d -> b c d h w", b=b, h=h, w=w) + + return x + + @property + def weight(self) -> torch.Tensor: + return self.weight2 + + +class CausalConv3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int = 3, + stride: Union[int, Tuple[int]] = 1, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, + ) -> None: + super().__init__() + + self.in_channels = in_channels + self.out_channels = out_channels + + kernel_size = (kernel_size, kernel_size, kernel_size) + self.time_kernel_size = kernel_size[0] + + dilation = (dilation, 1, 1) + + height_pad = kernel_size[1] // 2 + width_pad = kernel_size[2] // 2 + padding = (0, height_pad, width_pad) + + self.conv = nn.Conv3d( + in_channels, + out_channels, + kernel_size, + stride=stride, + dilation=dilation, + padding=padding, + padding_mode=spatial_padding_mode.value, + groups=groups, + bias=bias, + ) + + def forward(self, x: torch.Tensor, causal: bool = True) -> torch.Tensor: + if causal: + first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, self.time_kernel_size - 1, 1, 1)) + x = torch.concatenate((first_frame_pad, x), dim=2) + else: + first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1)) + last_frame_pad = x[:, :, -1:, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1)) + x = torch.concatenate((first_frame_pad, x, last_frame_pad), dim=2) + x = self.conv(x) + return x + + @property + def weight(self) -> torch.Tensor: + return self.conv.weight + + +from torch import nn + + +def patchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: + """ + Rearrange spatial dimensions into channels. Divides image into patch_size x patch_size blocks + and moves pixels from each block into separate channels (space-to-depth). + Args: + x: Input tensor (4D or 5D) + patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, divides HxW into 4x4 blocks. + patch_size_t: Temporal patch size for frames. Default=1 (no temporal patching). + For 5D: (B, C, F, H, W) -> (B, Cx(patch_size_hw^2)x(patch_size_t), F/patch_size_t, H/patch_size_hw, W/patch_size_hw) + Example: (B, 3, 33, 512, 512) with patch_size_hw=4, patch_size_t=1 -> (B, 48, 33, 128, 128) + """ + if patch_size_hw == 1 and patch_size_t == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size_hw, r=patch_size_hw) + elif x.dim() == 5: + x = rearrange( + x, + "b c (f p) (h q) (w r) -> b (c p r q) f h w", + p=patch_size_t, + q=patch_size_hw, + r=patch_size_hw, + ) + else: + raise ValueError(f"Invalid input shape: {x.shape}") + + return x + + +def unpatchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: + """ + Rearrange channels back into spatial dimensions. Inverse of patchify - moves pixels from + channels back into patch_size x patch_size blocks (depth-to-space). + Args: + x: Input tensor (4D or 5D) + patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, expands HxW by 4x. + patch_size_t: Temporal patch size for frames. Default=1 (no temporal expansion). + For 5D: (B, Cx(patch_size_hw^2)x(patch_size_t), F, H, W) -> (B, C, Fxpatch_size_t, Hxpatch_size_hw, Wxpatch_size_hw) + Example: (B, 48, 33, 128, 128) with patch_size_hw=4, patch_size_t=1 -> (B, 3, 33, 512, 512) + """ + if patch_size_hw == 1 and patch_size_t == 1: + return x + + if x.dim() == 4: + x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size_hw, r=patch_size_hw) + elif x.dim() == 5: + x = rearrange( + x, + "b (c p r q) f h w -> b c (f p) (h q) (w r)", + p=patch_size_t, + q=patch_size_hw, + r=patch_size_hw, + ) + + return x + + +class PerChannelStatistics(nn.Module): + """ + Per-channel statistics for normalizing and denormalizing the latent representation. + This statics is computed over the entire dataset and stored in model's checkpoint under VAE state_dict. + """ + + def __init__(self, latent_channels: int = 128): + super().__init__() + self.register_buffer("std-of-means", torch.empty(latent_channels)) + self.register_buffer("mean-of-means", torch.empty(latent_channels)) + + def un_normalize(self, x: torch.Tensor) -> torch.Tensor: + return (x * self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)) + self.get_buffer("mean-of-means").view( + 1, -1, 1, 1, 1 + ).to(x) + + def normalize(self, x: torch.Tensor) -> torch.Tensor: + return (x - self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)) / self.get_buffer("std-of-means").view( + 1, -1, 1, 1, 1 + ).to(x) + + +import math + +from torch import nn + + +class SpaceToDepthDownsample(nn.Module): + def __init__( + self, + dims: Union[int, Tuple[int, int]], + in_channels: int, + out_channels: int, + stride: Tuple[int, int, int], + spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, + ): + super().__init__() + self.stride = stride + self.group_size = in_channels * math.prod(stride) // out_channels + self.conv = make_conv_nd( + dims=dims, + in_channels=in_channels, + out_channels=out_channels // math.prod(stride), + kernel_size=3, + stride=1, + causal=True, + spatial_padding_mode=spatial_padding_mode, + ) + + def forward( + self, + x: torch.Tensor, + causal: bool = True, + ) -> torch.Tensor: + if self.stride[0] == 2: + x = torch.cat([x[:, :, :1, :, :], x], dim=2) # duplicate first frames for padding + + # skip connection + x_in = rearrange( + x, + "b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w", + p1=self.stride[0], + p2=self.stride[1], + p3=self.stride[2], + ) + x_in = rearrange(x_in, "b (c g) d h w -> b c g d h w", g=self.group_size) + x_in = x_in.mean(dim=2) + + # conv + x = self.conv(x, causal=causal) + x = rearrange( + x, + "b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w", + p1=self.stride[0], + p2=self.stride[1], + p3=self.stride[2], + ) + + x = x + x_in + + return x + + +class DepthToSpaceUpsample(nn.Module): + def __init__( + self, + dims: int | Tuple[int, int], + in_channels: int, + stride: Tuple[int, int, int], + residual: bool = False, + out_channels_reduction_factor: int = 1, + spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, + ): + super().__init__() + self.stride = stride + self.out_channels = math.prod(stride) * in_channels // out_channels_reduction_factor + self.conv = make_conv_nd( + dims=dims, + in_channels=in_channels, + out_channels=self.out_channels, + kernel_size=3, + stride=1, + causal=True, + spatial_padding_mode=spatial_padding_mode, + ) + self.residual = residual + self.out_channels_reduction_factor = out_channels_reduction_factor + + def forward( + self, + x: torch.Tensor, + causal: bool = True, + ) -> torch.Tensor: + if self.residual: + # Reshape and duplicate the input to match the output shape + x_in = rearrange( + x, + "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)", + p1=self.stride[0], + p2=self.stride[1], + p3=self.stride[2], + ) + num_repeat = math.prod(self.stride) // self.out_channels_reduction_factor + x_in = x_in.repeat(1, num_repeat, 1, 1, 1) + if self.stride[0] == 2: + x_in = x_in[:, :, 1:, :, :] + x = self.conv(x, causal=causal) + x = rearrange( + x, + "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)", + p1=self.stride[0], + p2=self.stride[1], + p3=self.stride[2], + ) + if self.stride[0] == 2: + x = x[:, :, 1:, :, :] + if self.residual: + x = x + x_in + return x + + +import itertools +from dataclasses import dataclass +from typing import Callable, List, NamedTuple + + +def compute_trapezoidal_mask_1d( + length: int, + ramp_left: int, + ramp_right: int, + left_starts_from_0: bool = False, +) -> torch.Tensor: + """ + Generate a 1D trapezoidal blending mask with linear ramps. + Args: + length: Output length of the mask. + ramp_left: Fade-in length on the left. + ramp_right: Fade-out length on the right. + left_starts_from_0: Whether the ramp starts from 0 or first non-zero value. + Useful for temporal tiles where the first tile is causal. + Returns: + A 1D tensor of shape `(length,)` with values in [0, 1]. + """ + if length <= 0: + raise ValueError("Mask length must be positive.") + + ramp_left = max(0, min(ramp_left, length)) + ramp_right = max(0, min(ramp_right, length)) + + mask = torch.ones(length) + + if ramp_left > 0: + interval_length = ramp_left + 1 if left_starts_from_0 else ramp_left + 2 + fade_in = torch.linspace(0.0, 1.0, interval_length)[:-1] + if not left_starts_from_0: + fade_in = fade_in[1:] + mask[:ramp_left] *= fade_in + + if ramp_right > 0: + fade_out = torch.linspace(1.0, 0.0, steps=ramp_right + 2)[1:-1] + mask[-ramp_right:] *= fade_out + + return mask.clamp_(0, 1) + + +def compute_rectangular_mask_1d( + length: int, + left_ramp: int, + right_ramp: int, +) -> torch.Tensor: + """ + Generate a 1D rectangular (pulse) mask. + Args: + length: Output length of the mask. + left_ramp: Number of elements at the start of the mask to set to 0. + right_ramp: Number of elements at the end of the mask to set to 0. + Returns: + A 1D tensor of shape `(length,)` with values 0 or 1. + """ + if length <= 0: + raise ValueError("Mask length must be positive.") + + mask = torch.ones(length) + if left_ramp > 0: + mask[:left_ramp] = 0 + if right_ramp > 0: + mask[-right_ramp:] = 0 + return mask + + +def _validate_spatial_tiling(tile_size_in_pixels: int, tile_overlap_in_pixels: int) -> None: + """Validate spatial tiling arguments for the LTX VAE.""" + if tile_size_in_pixels < 64: + raise ValueError(f"tile_size_in_pixels must be at least 64, got {tile_size_in_pixels}") + if tile_size_in_pixels % 32 != 0: + raise ValueError(f"tile_size_in_pixels must be divisible by 32, got {tile_size_in_pixels}") + if tile_overlap_in_pixels % 32 != 0: + raise ValueError(f"tile_overlap_in_pixels must be divisible by 32, got {tile_overlap_in_pixels}") + if tile_overlap_in_pixels >= tile_size_in_pixels: + raise ValueError(f"Overlap must be less than tile size, got {tile_overlap_in_pixels} and {tile_size_in_pixels}") + + +def _validate_temporal_tiling(tile_size_in_frames: int, tile_overlap_in_frames: int) -> None: + """Validate temporal tiling arguments for the LTX VAE.""" + if tile_size_in_frames < 16: + raise ValueError(f"tile_size_in_frames must be at least 16, got {tile_size_in_frames}") + if tile_size_in_frames % 8 != 0: + raise ValueError(f"tile_size_in_frames must be divisible by 8, got {tile_size_in_frames}") + if tile_overlap_in_frames % 8 != 0: + raise ValueError(f"tile_overlap_in_frames must be divisible by 8, got {tile_overlap_in_frames}") + if tile_overlap_in_frames >= tile_size_in_frames: + raise ValueError(f"Overlap must be less than tile size, got {tile_overlap_in_frames} and {tile_size_in_frames}") + + +@dataclass(frozen=True) +class DimensionIntervals: + """Defines how a single dimension is split into overlapping intervals (tiles). + Each list has length N where N is the number of intervals. The i-th element + of each list describes the i-th interval. + Attributes: + starts: Start index of each interval (inclusive). + ends: End index of each interval (exclusive). + left_ramps: Length of the left blend ramp for each interval. + Used to create masks that fade in from 0 to 1. + right_ramps: Length of the right blend ramp for each interval. + Used to create masks that fade out from 1 to 0. + """ + + starts: List[int] + ends: List[int] + left_ramps: List[int] + right_ramps: List[int] + + +# Operation to split a single dimension of the tensor into intervals based on the length along the dimension. +SplitOperation = Callable[[int], DimensionIntervals] +# Operation to map the intervals in input dimension to slices and masks along a corresponding output dimension. +MappingOperation = Callable[[DimensionIntervals], tuple[list[slice], list[torch.Tensor | None]]] + + +def default_split_operation(length: int) -> DimensionIntervals: + return DimensionIntervals(starts=[0], ends=[length], left_ramps=[0], right_ramps=[0]) + + +DEFAULT_SPLIT_OPERATION: SplitOperation = default_split_operation + + +def default_mapping_operation( + _intervals: DimensionIntervals, +) -> tuple[list[slice], list[torch.Tensor | None]]: + return [slice(0, None)], [None] + + +DEFAULT_MAPPING_OPERATION: MappingOperation = default_mapping_operation + + +class Tile(NamedTuple): + """ + Represents a single tile. + Attributes: + in_coords: + Tuple of slices specifying where to cut the tile from the INPUT tensor. + out_coords: + Tuple of slices specifying where this tile's OUTPUT should be placed in the reconstructed OUTPUT tensor. + masks_1d: + Per-dimension masks in OUTPUT units. + These are used to create all-dimensional blending mask. + Methods: + blend_mask: + Create a single N-D mask from the per-dimension masks. + """ + + in_coords: Tuple[slice, ...] + out_coords: Tuple[slice, ...] + masks_1d: Tuple[Tuple[torch.Tensor, ...]] + + @property + def blend_mask(self) -> torch.Tensor: + num_dims = len(self.out_coords) + per_dimension_masks: List[torch.Tensor] = [] + + for dim_idx in range(num_dims): + mask_1d = self.masks_1d[dim_idx] + view_shape = [1] * num_dims + if mask_1d is None: + # Broadcast mask along this dimension (length 1). + one = torch.ones(1) + + view_shape[dim_idx] = 1 + per_dimension_masks.append(one.view(*view_shape)) + continue + + # Reshape (L,) -> (1, ..., L, ..., 1) so masks across dimensions broadcast-multiply. + view_shape[dim_idx] = mask_1d.shape[0] + per_dimension_masks.append(mask_1d.view(*view_shape)) + + # Multiply per-dimension masks to form the full N-D mask (separable blending window). + combined_mask = per_dimension_masks[0] + for mask in per_dimension_masks[1:]: + combined_mask = combined_mask * mask + + return combined_mask + + +def create_tiles_from_intervals_and_mappers( + original_shape: torch.Size, + dimension_intervals: Tuple[DimensionIntervals, ...], + mappers: List[MappingOperation], +) -> List[Tile]: + full_dim_input_slices = [] + full_dim_output_slices = [] + full_dim_masks_1d = [] + for axis_index in range(len(original_shape)): + intervals = dimension_intervals[axis_index] + starts = intervals.starts + ends = intervals.ends + input_slices = [slice(s, e) for s, e in zip(starts, ends, strict=True)] + output_slices, masks_1d = mappers[axis_index](intervals) + full_dim_input_slices.append(input_slices) + full_dim_output_slices.append(output_slices) + full_dim_masks_1d.append(masks_1d) + + tiles = [] + tile_in_coords = list(itertools.product(*full_dim_input_slices)) + tile_out_coords = list(itertools.product(*full_dim_output_slices)) + tile_mask_1ds = list(itertools.product(*full_dim_masks_1d)) + for in_coord, out_coord, mask_1d in zip(tile_in_coords, tile_out_coords, tile_mask_1ds, strict=True): + tiles.append( + Tile( + in_coords=in_coord, + out_coords=out_coord, + masks_1d=mask_1d, + ) + ) + return tiles + + +def create_tiles( + tensor_shape: torch.Size, + splitters: List[SplitOperation], + mappers: List[MappingOperation], +) -> List[Tile]: + if len(splitters) != len(tensor_shape): + raise ValueError( + f"Number of splitters must be equal to number of dimensions in tensor shape, " + f"got {len(splitters)} and {len(tensor_shape)}" + ) + if len(mappers) != len(tensor_shape): + raise ValueError( + f"Number of mappers must be equal to number of dimensions in tensor shape, " + f"got {len(mappers)} and {len(tensor_shape)}" + ) + dimension_intervals = tuple(splitter(length) for splitter, length in zip(splitters, tensor_shape, strict=True)) + return create_tiles_from_intervals_and_mappers(tensor_shape, dimension_intervals, mappers) + + +from typing import Optional + +import torch +from torch import nn + +from .transformer import PixArtAlphaCombinedTimestepSizeEmbeddings + + +class ResnetBlock3D(nn.Module): + r""" + A Resnet block. + Parameters: + in_channels (`int`): The number of channels in the input. + out_channels (`int`, *optional*, default to be `None`): + The number of output channels for the first conv layer. If None, same as `in_channels`. + dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use. + groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer. + eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization. + """ + + def __init__( + self, + dims: Union[int, Tuple[int, int]], + in_channels: int, + out_channels: Optional[int] = None, + dropout: float = 0.0, + groups: int = 32, + eps: float = 1e-6, + norm_layer: NormLayerType = NormLayerType.PIXEL_NORM, + inject_noise: bool = False, + timestep_conditioning: bool = False, + spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, + ): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.inject_noise = inject_noise + + if norm_layer == NormLayerType.GROUP_NORM: + self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True) + elif norm_layer == NormLayerType.PIXEL_NORM: + self.norm1 = PixelNorm() + + self.non_linearity = nn.SiLU() + + self.conv1 = make_conv_nd( + dims, + in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1, + causal=True, + spatial_padding_mode=spatial_padding_mode, + ) + + if inject_noise: + self.per_channel_scale1 = nn.Parameter(torch.zeros((in_channels, 1, 1))) + + if norm_layer == NormLayerType.GROUP_NORM: + self.norm2 = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps, affine=True) + elif norm_layer == NormLayerType.PIXEL_NORM: + self.norm2 = PixelNorm() + + self.dropout = torch.nn.Dropout(dropout) + + self.conv2 = make_conv_nd( + dims, + out_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1, + causal=True, + spatial_padding_mode=spatial_padding_mode, + ) + + if inject_noise: + self.per_channel_scale2 = nn.Parameter(torch.zeros((in_channels, 1, 1))) + + self.conv_shortcut = ( + make_linear_nd(dims=dims, in_channels=in_channels, out_channels=out_channels) + if in_channels != out_channels + else nn.Identity() + ) + + # Using GroupNorm with 1 group is equivalent to LayerNorm but works with (B, C, ...) layout + # avoiding the need for dimension rearrangement used in standard nn.LayerNorm + self.norm3 = ( + nn.GroupNorm(num_groups=1, num_channels=in_channels, eps=eps, affine=True) + if in_channels != out_channels + else nn.Identity() + ) + + self.timestep_conditioning = timestep_conditioning + + if timestep_conditioning: + self.scale_shift_table = nn.Parameter(torch.zeros(4, in_channels)) + + def _feed_spatial_noise( + self, + hidden_states: torch.Tensor, + per_channel_scale: torch.Tensor, + generator: Optional[torch.Generator] = None, + ) -> torch.Tensor: + spatial_shape = hidden_states.shape[-2:] + device = hidden_states.device + dtype = hidden_states.dtype + + # similar to the "explicit noise inputs" method in style-gan + spatial_noise = torch.randn(spatial_shape, device=device, dtype=dtype, generator=generator)[None] + scaled_noise = (spatial_noise * per_channel_scale)[None, :, None, ...] + hidden_states = hidden_states + scaled_noise + + return hidden_states + + def forward( + self, + input_tensor: torch.Tensor, + causal: bool = True, + timestep: Optional[torch.Tensor] = None, + generator: Optional[torch.Generator] = None, + ) -> torch.Tensor: + hidden_states = input_tensor + batch_size = hidden_states.shape[0] + + hidden_states = self.norm1(hidden_states) + if self.timestep_conditioning: + if timestep is None: + raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True") + ada_values = self.scale_shift_table[None, ..., None, None, None].to( + device=hidden_states.device, dtype=hidden_states.dtype + ) + timestep.reshape( + batch_size, + 4, + -1, + timestep.shape[-3], + timestep.shape[-2], + timestep.shape[-1], + ) + shift1, scale1, shift2, scale2 = ada_values.unbind(dim=1) + + hidden_states = hidden_states * (1 + scale1) + shift1 + + hidden_states = self.non_linearity(hidden_states) + + hidden_states = self.conv1(hidden_states, causal=causal) + + if self.inject_noise: + hidden_states = self._feed_spatial_noise( + hidden_states, + self.per_channel_scale1.to(device=hidden_states.device, dtype=hidden_states.dtype), + generator=generator, + ) + + hidden_states = self.norm2(hidden_states) + + if self.timestep_conditioning: + hidden_states = hidden_states * (1 + scale2) + shift2 + + hidden_states = self.non_linearity(hidden_states) + + hidden_states = self.dropout(hidden_states) + + hidden_states = self.conv2(hidden_states, causal=causal) + + if self.inject_noise: + hidden_states = self._feed_spatial_noise( + hidden_states, + self.per_channel_scale2.to(device=hidden_states.device, dtype=hidden_states.dtype), + generator=generator, + ) + + input_tensor = self.norm3(input_tensor) + + batch_size = input_tensor.shape[0] + + input_tensor = self.conv_shortcut(input_tensor) + + output_tensor = input_tensor + hidden_states + + return output_tensor + + +class UNetMidBlock3D(nn.Module): + """ + A 3D UNet mid-block [`UNetMidBlock3D`] with multiple residual blocks. + Args: + in_channels (`int`): The number of input channels. + dropout (`float`, *optional*, defaults to 0.0): The dropout rate. + num_layers (`int`, *optional*, defaults to 1): The number of residual blocks. + resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks. + resnet_groups (`int`, *optional*, defaults to 32): + The number of groups to use in the group normalization layers of the resnet blocks. + norm_layer (`str`, *optional*, defaults to `group_norm`): + The normalization layer to use. Can be either `group_norm` or `pixel_norm`. + inject_noise (`bool`, *optional*, defaults to `False`): + Whether to inject noise into the hidden states. + timestep_conditioning (`bool`, *optional*, defaults to `False`): + Whether to condition the hidden states on the timestep. + Returns: + `torch.Tensor`: The output of the last residual block, which is a tensor of shape `(batch_size, + in_channels, height, width)`. + """ + + def __init__( + self, + dims: Union[int, Tuple[int, int]], + in_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_groups: int = 32, + norm_layer: NormLayerType = NormLayerType.GROUP_NORM, + inject_noise: bool = False, + timestep_conditioning: bool = False, + spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, + ): + super().__init__() + resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + + self.timestep_conditioning = timestep_conditioning + + if timestep_conditioning: + self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings( + embedding_dim=in_channels * 4, size_emb_dim=0 + ) + + self.res_blocks = nn.ModuleList( + [ + ResnetBlock3D( + dims=dims, + in_channels=in_channels, + out_channels=in_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + norm_layer=norm_layer, + inject_noise=inject_noise, + timestep_conditioning=timestep_conditioning, + spatial_padding_mode=spatial_padding_mode, + ) + for _ in range(num_layers) + ] + ) + + def forward( + self, + hidden_states: torch.Tensor, + causal: bool = True, + timestep: Optional[torch.Tensor] = None, + generator: Optional[torch.Generator] = None, + ) -> torch.Tensor: + timestep_embed = None + if self.timestep_conditioning: + if timestep is None: + raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True") + batch_size = hidden_states.shape[0] + timestep_embed = self.time_embedder( + timestep=timestep.flatten(), + hidden_dtype=hidden_states.dtype, + ) + timestep_embed = timestep_embed.view(batch_size, timestep_embed.shape[-1], 1, 1, 1) + + for resnet in self.res_blocks: + hidden_states = resnet( + hidden_states, + causal=causal, + timestep=timestep_embed, + generator=generator, + ) + + return hidden_states + + +import logging +from dataclasses import replace +from typing import Any, Callable, Iterator, NamedTuple + +import torch +from torch import nn + +logger: logging.Logger = logging.getLogger(__name__) + + +class VideoPixelShape(NamedTuple): + batch: int + frames: int + width: int + height: int + fps: float + + def to_torch_shape(self, channels: int = 3) -> torch.Size: + return torch.Size([self.batch, channels, self.frames, self.height, self.width]) + + +class SpatioTemporalScaleFactors(NamedTuple): + time: int + width: int + height: int + + @classmethod + def default(cls) -> "SpatioTemporalScaleFactors": + return cls(time=8, width=32, height=32) + + +VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default() + + +class VideoLatentShape(NamedTuple): + batch: int + channels: int + frames: int + height: int + width: int + + def to_torch_shape(self) -> torch.Size: + return torch.Size([self.batch, self.channels, self.frames, self.height, self.width]) + + @staticmethod + def from_torch_shape(shape: torch.Size) -> "VideoLatentShape": + return VideoLatentShape(shape[0], shape[1], shape[2], shape[3], shape[4]) + + def token_count(self) -> int: + return self.frames * self.height * self.width + + def mask_shape(self) -> "VideoLatentShape": + return self._replace(channels=1) + + def upscale(self, scale_factors: SpatioTemporalScaleFactors) -> VideoPixelShape: + return VideoPixelShape( + batch=self.batch, + frames=(self.frames - 1) * scale_factors.time + 1, + width=self.width * scale_factors.width, + height=self.height * scale_factors.height, + fps=1.0, + ) + + +def _make_encoder_block( + block_name: str, + block_config: dict[str, Any], + in_channels: int, + convolution_dimensions: int, + norm_layer: NormLayerType, + norm_num_groups: int, + spatial_padding_mode: PaddingModeType, +) -> Tuple[nn.Module, int]: + out_channels = in_channels + + if block_name == "res_x": + block = UNetMidBlock3D( + dims=convolution_dimensions, + in_channels=in_channels, + num_layers=block_config["num_layers"], + resnet_eps=1e-6, + resnet_groups=norm_num_groups, + norm_layer=norm_layer, + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "res_x_y": + out_channels = in_channels * block_config.get("multiplier", 2) + block = ResnetBlock3D( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + eps=1e-6, + groups=norm_num_groups, + norm_layer=norm_layer, + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_time": + block = make_conv_nd( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=(2, 1, 1), + causal=True, + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_space": + block = make_conv_nd( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=(1, 2, 2), + causal=True, + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_all": + block = make_conv_nd( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=(2, 2, 2), + causal=True, + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_all_x_y": + out_channels = in_channels * block_config.get("multiplier", 2) + block = make_conv_nd( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=(2, 2, 2), + causal=True, + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_all_res": + out_channels = in_channels * block_config.get("multiplier", 2) + block = SpaceToDepthDownsample( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + stride=(2, 2, 2), + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_space_res": + out_channels = in_channels * block_config.get("multiplier", 2) + block = SpaceToDepthDownsample( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + stride=(1, 2, 2), + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_time_res": + out_channels = in_channels * block_config.get("multiplier", 2) + block = SpaceToDepthDownsample( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + stride=(2, 1, 1), + spatial_padding_mode=spatial_padding_mode, + ) + else: + raise ValueError(f"unknown block: {block_name}") + + return block, out_channels + + +class VideoEncoder(nn.Module): + _DEFAULT_NORM_NUM_GROUPS = 32 + """ + Variational Autoencoder Encoder. Encodes video frames into a latent representation. + The encoder compresses the input video through a series of downsampling operations controlled by + patch_size and encoder_blocks. The output is a normalized latent tensor with shape (B, 128, F', H', W'). + Compression Behavior: + The total compression is determined by: + 1. Initial spatial compression via patchify: H -> H/4, W -> W/4 (patch_size=4) + 2. Sequential compression through encoder_blocks based on their stride patterns + Compression blocks apply 2x compression in specified dimensions: + - "compress_time" / "compress_time_res": temporal only + - "compress_space" / "compress_space_res": spatial only (H and W) + - "compress_all" / "compress_all_res": all dimensions (F, H, W) + - "res_x" / "res_x_y": no compression + Standard LTX Video configuration: + - patch_size=4 + - encoder_blocks: 1x compress_space_res, 1x compress_time_res, 2x compress_all_res + - Final dimensions: F' = 1 + (F-1)/8, H' = H/32, W' = W/32 + - Example: (B, 3, 33, 512, 512) -> (B, 128, 5, 16, 16) + - Note: Input must have 1 + 8*k frames (e.g., 1, 9, 17, 25, 33...) + Args: + convolution_dimensions: The number of dimensions to use in convolutions (2D or 3D). + in_channels: The number of input channels. For RGB images, this is 3. + out_channels: The number of output channels (latent channels). For latent channels, this is 128. + encoder_blocks: The list of blocks to construct the encoder. Each block is a tuple of (block_name, params) + where params is either an int (num_layers) or a dict with configuration. + patch_size: The patch size for initial spatial compression. Should be a power of 2. + norm_layer: The normalization layer to use. Can be either `group_norm` or `pixel_norm`. + latent_log_var: The log variance mode. Can be either `per_channel`, `uniform`, `constant` or `none`. + """ + + def __init__( + self, + convolution_dimensions: int = 3, + in_channels: int = 3, + out_channels: int = 128, + encoder_blocks: List[Tuple[str, int]] | List[Tuple[str, dict[str, Any]]] = [], # noqa: B006 + patch_size: int = 4, + norm_layer: NormLayerType = NormLayerType.PIXEL_NORM, + latent_log_var: LogVarianceType = LogVarianceType.UNIFORM, + encoder_spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, + ): + super().__init__() + + self.patch_size = patch_size + self.norm_layer = norm_layer + self.latent_channels = out_channels + self.latent_log_var = latent_log_var + self._norm_num_groups = self._DEFAULT_NORM_NUM_GROUPS + + # Per-channel statistics for normalizing latents + self.per_channel_statistics = PerChannelStatistics(latent_channels=out_channels) + + in_channels = in_channels * patch_size**2 + feature_channels = out_channels + + self.conv_in = make_conv_nd( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=feature_channels, + kernel_size=3, + stride=1, + padding=1, + causal=True, + spatial_padding_mode=encoder_spatial_padding_mode, + ) + + self.down_blocks = nn.ModuleList([]) + + for block_name, block_params in encoder_blocks: + # Convert int to dict format for uniform handling + block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params + + block, feature_channels = _make_encoder_block( + block_name=block_name, + block_config=block_config, + in_channels=feature_channels, + convolution_dimensions=convolution_dimensions, + norm_layer=norm_layer, + norm_num_groups=self._norm_num_groups, + spatial_padding_mode=encoder_spatial_padding_mode, + ) + + self.down_blocks.append(block) + + # out + if norm_layer == NormLayerType.GROUP_NORM: + self.conv_norm_out = nn.GroupNorm(num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6) + elif norm_layer == NormLayerType.PIXEL_NORM: + self.conv_norm_out = PixelNorm() + + self.conv_act = nn.SiLU() + + conv_out_channels = out_channels + if latent_log_var == LogVarianceType.PER_CHANNEL: + conv_out_channels *= 2 + elif latent_log_var in {LogVarianceType.UNIFORM, LogVarianceType.CONSTANT}: + conv_out_channels += 1 + elif latent_log_var != LogVarianceType.NONE: + raise ValueError(f"Invalid latent_log_var: {latent_log_var}") + + self.conv_out = make_conv_nd( + dims=convolution_dimensions, + in_channels=feature_channels, + out_channels=conv_out_channels, + kernel_size=3, + padding=1, + causal=True, + spatial_padding_mode=encoder_spatial_padding_mode, + ) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + r""" + Encode video frames into normalized latent representation. + Args: + sample: Input video (B, C, F, H, W). F should be 1 + 8*k (e.g., 1, 9, 17, 25, 33...). + If not, the encoder crops the last frames to the nearest valid length. + Returns: + Normalized latent means (B, 128, F', H', W') where F' = 1+(F-1)/8, H' = H/32, W' = W/32. + Example: (B, 3, 33, 512, 512) -> (B, 128, 5, 16, 16). + """ + # Validate frame count (crop to nearest valid length if needed) + frames_count = sample.shape[2] + if ((frames_count - 1) % 8) != 0: + frames_to_crop = (frames_count - 1) % 8 + logger.warning( + "Invalid number of frames %s for encode; cropping last %s frames to satisfy 1 + 8*k.", + frames_count, + frames_to_crop, + ) + sample = sample[:, :, :-frames_to_crop, ...] + + # Initial spatial compression: trade spatial resolution for channel depth + # This reduces H,W by patch_size and increases channels, making convolutions more efficient + # Example: (B, 3, F, 512, 512) -> (B, 48, F, 128, 128) with patch_size=4 + sample = patchify(sample, patch_size_hw=self.patch_size, patch_size_t=1) + sample = self.conv_in(sample) + + for down_block in self.down_blocks: + sample = down_block(sample) + + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + if self.latent_log_var == LogVarianceType.UNIFORM: + # Uniform Variance: model outputs N means and 1 shared log-variance channel. + # We need to expand the single logvar to match the number of means channels + # to create a format compatible with PER_CHANNEL (means + logvar, each with N channels). + # Sample shape: (B, N+1, ...) where N = latent_channels (e.g., 128 means + 1 logvar = 129) + # Target shape: (B, 2*N, ...) where first N are means, last N are logvar + + if sample.shape[1] < 2: + raise ValueError( + f"Invalid channel count for UNIFORM mode: expected at least 2 channels " + f"(N means + 1 logvar), got {sample.shape[1]}" + ) + + # Extract means (first N channels) and logvar (last 1 channel) + means = sample[:, :-1, ...] # (B, N, ...) + logvar = sample[:, -1:, ...] # (B, 1, ...) + + # Repeat logvar N times to match means channels + # Use expand/repeat pattern that works for both 4D and 5D tensors + num_channels = means.shape[1] + repeat_shape = [1, num_channels] + [1] * (sample.ndim - 2) + repeated_logvar = logvar.repeat(*repeat_shape) # (B, N, ...) + + # Concatenate to create (B, 2*N, ...) format: [means, repeated_logvar] + sample = torch.cat([means, repeated_logvar], dim=1) + elif self.latent_log_var == LogVarianceType.CONSTANT: + sample = sample[:, :-1, ...] + approx_ln_0 = -30 # this is the minimal clamp value in DiagonalGaussianDistribution objects + sample = torch.cat( + [sample, torch.ones_like(sample, device=sample.device) * approx_ln_0], + dim=1, + ) + + # Split into means and logvar, then normalize means + means, _ = torch.chunk(sample, 2, dim=1) + return self.per_channel_statistics.normalize(means) + + def tiled_encode( + self, + video: torch.Tensor, + tile_size: tuple[int, int] = (64, 512), + tile_stride: tuple[int, int] = (40, 448), + ) -> torch.Tensor: + """Encode video to latent using tiled processing of the given video tensor. + Device Handling: + - Input video can be on CPU or GPU + - Accumulation buffers are created on model's device + - Each tile is automatically moved to model's device before encoding + - Output latent is returned on model's device + Args: + video: Input video tensor (B, 3, F, H, W) in range [-1, 1] + tile_size: Temporal/spatial tile size as `(frames, pixels)`. + tile_stride: Temporal/spatial tile stride as `(frames, pixels)`. + Returns: + Latent tensor (B, 128, F', H', W') on model's device + where F' = 1 + (F-1)/8, H' = H/32, W' = W/32 + """ + # Detect model device and dtype + model_device = next(self.parameters()).device + model_dtype = next(self.parameters()).dtype + + # Extract shape components + batch, _, frames, height, width = video.shape + + # Check frame count and crop if needed + if (frames - 1) % VIDEO_SCALE_FACTORS.time != 0: + frames_to_crop = (frames - 1) % VIDEO_SCALE_FACTORS.time + logger.warning( + f"Number of frames {frames} of input video is not ({VIDEO_SCALE_FACTORS.time} * k + 1), " + f"last {frames_to_crop} frames will be cropped" + ) + video = video[:, :, :-frames_to_crop, ...] + # Update frames after cropping + frames = video.shape[2] + + # Calculate output latent shape (inverse of upscale) + latent_shape = VideoLatentShape( + batch=batch, + channels=self.latent_channels, # 128 for standard VAE + frames=(frames - 1) // VIDEO_SCALE_FACTORS.time + 1, + height=height // VIDEO_SCALE_FACTORS.height, + width=width // VIDEO_SCALE_FACTORS.width, + ) + + # Prepare tiles (operates on VIDEO dimensions) + temporal_tile_size, spatial_tile_size = tile_size + temporal_tile_stride, spatial_tile_stride = tile_stride + tiles = prepare_tiles_for_encoding( + video, + spatial_tile_size_in_pixels=spatial_tile_size, + spatial_tile_overlap_in_pixels=spatial_tile_size - spatial_tile_stride, + temporal_tile_size_in_frames=temporal_tile_size, + temporal_tile_overlap_in_frames=temporal_tile_size - temporal_tile_stride, + ) + + # Initialize accumulation buffers on model device + latent_buffer = torch.zeros( + latent_shape.to_torch_shape(), + device=model_device, + dtype=model_dtype, + ) + weights_buffer = torch.zeros_like(latent_buffer) + + # Process each tile + for tile in tiles: + # Extract video tile from input (may be on CPU) + video_tile = video[tile.in_coords] + + # Move tile to model device if needed + if video_tile.device != model_device or video_tile.dtype != model_dtype: + video_tile = video_tile.to(device=model_device, dtype=model_dtype) + + # Encode tile to latent (output on model device) + latent_tile = self.forward(video_tile) + + # Move blend mask to model device + mask = tile.blend_mask.to( + device=model_device, + dtype=model_dtype, + ) + + # Weighted accumulation in latent space + latent_buffer[tile.out_coords] += latent_tile * mask + weights_buffer[tile.out_coords] += mask + + del latent_tile, mask, video_tile + + # Normalize by accumulated weights + weights_buffer = weights_buffer.clamp(min=1e-8) + return latent_buffer / weights_buffer + + +def prepare_tiles_for_encoding( + video: torch.Tensor, + spatial_tile_size_in_pixels: int | None = None, + spatial_tile_overlap_in_pixels: int = 0, + temporal_tile_size_in_frames: int | None = None, + temporal_tile_overlap_in_frames: int = 0, +) -> List[Tile]: + """Prepare tiles for VAE encoding. + Args: + video: Input video tensor (B, 3, F, H, W) in range [-1, 1] + spatial_tile_size_in_pixels: Spatial tile size in pixels. + spatial_tile_overlap_in_pixels: Spatial overlap in pixels. + temporal_tile_size_in_frames: Temporal tile size in frames. + temporal_tile_overlap_in_frames: Temporal overlap in frames. + Returns: + List of tiles for the video tensor + """ + + splitters = [DEFAULT_SPLIT_OPERATION] * len(video.shape) + mappers = [DEFAULT_MAPPING_OPERATION] * len(video.shape) + minimum_spatial_overlap_px = 64 + minimum_temporal_overlap_frames = 16 + + if spatial_tile_size_in_pixels is not None: + _validate_spatial_tiling(spatial_tile_size_in_pixels, spatial_tile_overlap_in_pixels) + tile_size_px = spatial_tile_size_in_pixels + overlap_px = spatial_tile_overlap_in_pixels + + # Set minimum spatial overlap to 64 pixels in order to allow cutting padding from + # the front and back of the tiles and concatenate tiles without artifacts. + # The encoder uses symmetric padding (pad=1) in H and W at each conv layer. At tile + # boundaries, convs see padding (zeros/reflect) instead of real neighbor pixels, causing + # incorrect context near edges. + # For each overlap we discard 1 latent per edge (32px at scale 32) and concatenate tiles at a + # shared region with the next tile. + if overlap_px < minimum_spatial_overlap_px: + logger.warning( + f"Overlap pixels {overlap_px} in spatial tiling is less than \ + {minimum_spatial_overlap_px}, setting to minimum required {minimum_spatial_overlap_px}" + ) + overlap_px = minimum_spatial_overlap_px + + # Define split and map operations for the spatial dimensions + + # Height axis (H) + splitters[3] = split_with_symmetric_overlaps(tile_size_px, overlap_px) + mappers[3] = make_mapping_operation(map_spatial_interval_to_latent, scale=VIDEO_SCALE_FACTORS.height) + + # Width axis (W) + splitters[4] = split_with_symmetric_overlaps(tile_size_px, overlap_px) + mappers[4] = make_mapping_operation(map_spatial_interval_to_latent, scale=VIDEO_SCALE_FACTORS.width) + + if temporal_tile_size_in_frames is not None: + _validate_temporal_tiling(temporal_tile_size_in_frames, temporal_tile_overlap_in_frames) + tile_size_frames = temporal_tile_size_in_frames + overlap_frames = temporal_tile_overlap_in_frames + + if overlap_frames < minimum_temporal_overlap_frames: + logger.warning(f"Overlap frames {overlap_frames} is less than 16, setting to minimum required 16") + overlap_frames = minimum_temporal_overlap_frames + + splitters[2] = split_temporal_frames(tile_size_frames, overlap_frames) + mappers[2] = make_mapping_operation(map_temporal_interval_to_latent, scale=VIDEO_SCALE_FACTORS.time) + + return create_tiles(video.shape, splitters, mappers) + + +def _make_decoder_block( + block_name: str, + block_config: dict[str, Any], + in_channels: int, + convolution_dimensions: int, + norm_layer: NormLayerType, + timestep_conditioning: bool, + norm_num_groups: int, + spatial_padding_mode: PaddingModeType, +) -> Tuple[nn.Module, int]: + out_channels = in_channels + if block_name == "res_x": + block = UNetMidBlock3D( + dims=convolution_dimensions, + in_channels=in_channels, + num_layers=block_config["num_layers"], + resnet_eps=1e-6, + resnet_groups=norm_num_groups, + norm_layer=norm_layer, + inject_noise=block_config.get("inject_noise", False), + timestep_conditioning=timestep_conditioning, + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "attn_res_x": + block = UNetMidBlock3D( + dims=convolution_dimensions, + in_channels=in_channels, + num_layers=block_config["num_layers"], + resnet_groups=norm_num_groups, + norm_layer=norm_layer, + inject_noise=block_config.get("inject_noise", False), + timestep_conditioning=timestep_conditioning, + attention_head_dim=block_config["attention_head_dim"], + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "res_x_y": + out_channels = in_channels // block_config.get("multiplier", 2) + block = ResnetBlock3D( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=out_channels, + eps=1e-6, + groups=norm_num_groups, + norm_layer=norm_layer, + inject_noise=block_config.get("inject_noise", False), + timestep_conditioning=False, + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_time": + out_channels = in_channels // block_config.get("multiplier", 1) + block = DepthToSpaceUpsample( + dims=convolution_dimensions, + in_channels=in_channels, + stride=(2, 1, 1), + out_channels_reduction_factor=block_config.get("multiplier", 1), + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_space": + out_channels = in_channels // block_config.get("multiplier", 1) + block = DepthToSpaceUpsample( + dims=convolution_dimensions, + in_channels=in_channels, + stride=(1, 2, 2), + out_channels_reduction_factor=block_config.get("multiplier", 1), + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_all": + out_channels = in_channels // block_config.get("multiplier", 1) + block = DepthToSpaceUpsample( + dims=convolution_dimensions, + in_channels=in_channels, + stride=(2, 2, 2), + residual=block_config.get("residual", False), + out_channels_reduction_factor=block_config.get("multiplier", 1), + spatial_padding_mode=spatial_padding_mode, + ) + else: + raise ValueError(f"unknown layer: {block_name}") + + return block, out_channels + + +class VideoDecoder(nn.Module): + _DEFAULT_NORM_NUM_GROUPS = 32 + """ + Variational Autoencoder Decoder. Decodes latent representation into video frames. + The decoder upsamples latents through a series of upsampling operations (inverse of encoder). + Output dimensions: F = 8x(F'-1) + 1, H = 32xH', W = 32xW' for standard LTX Video configuration. + Upsampling blocks expand dimensions by 2x in specified dimensions: + - "compress_time": temporal only + - "compress_space": spatial only (H and W) + - "compress_all": all dimensions (F, H, W) + - "res_x" / "res_x_y" / "attn_res_x": no upsampling + Causal Mode: + causal=False (standard): Symmetric padding, allows future frame dependencies. + causal=True: Causal padding, each frame depends only on past/current frames. + First frame removed after temporal upsampling in both modes. Output shape unchanged. + Example: (B, 128, 5, 16, 16) -> (B, 3, 33, 512, 512) for both modes. + Args: + convolution_dimensions: The number of dimensions to use in convolutions (2D or 3D). + in_channels: The number of input channels (latent channels). Default is 128. + out_channels: The number of output channels. For RGB images, this is 3. + decoder_blocks: The list of blocks to construct the decoder. Each block is a tuple of (block_name, params) + where params is either an int (num_layers) or a dict with configuration. + patch_size: Final spatial expansion factor. For standard LTX Video, use 4 for 4x spatial expansion: + H -> Hx4, W -> Wx4. Should be a power of 2. + norm_layer: The normalization layer to use. Can be either `group_norm` or `pixel_norm`. + causal: Whether to use causal convolutions. For standard LTX Video, use False for symmetric padding. + When True, uses causal padding (past/current frames only). + timestep_conditioning: Whether to condition the decoder on timestep for denoising. + """ + + def __init__( + self, + convolution_dimensions: int = 3, + in_channels: int = 128, + out_channels: int = 3, + decoder_blocks: List[Tuple[str, int | dict]] = [], # noqa: B006 + patch_size: int = 4, + norm_layer: NormLayerType = NormLayerType.PIXEL_NORM, + causal: bool = False, + timestep_conditioning: bool = False, + decoder_spatial_padding_mode: PaddingModeType = PaddingModeType.REFLECT, + base_channels: int = 128, + ): + super().__init__() + + # Spatiotemporal downscaling between decoded video space and VAE latents. + # According to the LTXV paper, the standard configuration downsamples + # video inputs by a factor of 8 in the temporal dimension and 32 in + # each spatial dimension (height and width). This parameter determines how + # many video frames and pixels correspond to a single latent cell. + self.video_downscale_factors = SpatioTemporalScaleFactors( + time=8, + width=32, + height=32, + ) + + self.patch_size = patch_size + out_channels = out_channels * patch_size**2 + self.causal = causal + self.timestep_conditioning = timestep_conditioning + self._norm_num_groups = self._DEFAULT_NORM_NUM_GROUPS + + # Per-channel statistics for denormalizing latents + self.per_channel_statistics = PerChannelStatistics(latent_channels=in_channels) + + # Noise and timestep parameters for decoder conditioning + self.decode_noise_scale = 0.025 + self.decode_timestep = 0.05 + + # LTX VAE decoder architecture uses 3 upsampler blocks with multiplier equals to 2. + # Hence the total feature_channels is multiplied by 8 (2^3). + feature_channels = base_channels * 8 + + self.conv_in = make_conv_nd( + dims=convolution_dimensions, + in_channels=in_channels, + out_channels=feature_channels, + kernel_size=3, + stride=1, + padding=1, + causal=True, + spatial_padding_mode=decoder_spatial_padding_mode, + ) + + self.up_blocks = nn.ModuleList([]) + + for block_name, block_params in list(reversed(decoder_blocks)): + # Convert int to dict format for uniform handling + block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params + + block, feature_channels = _make_decoder_block( + block_name=block_name, + block_config=block_config, + in_channels=feature_channels, + convolution_dimensions=convolution_dimensions, + norm_layer=norm_layer, + timestep_conditioning=timestep_conditioning, + norm_num_groups=self._norm_num_groups, + spatial_padding_mode=decoder_spatial_padding_mode, + ) + + self.up_blocks.append(block) + + if norm_layer == NormLayerType.GROUP_NORM: + self.conv_norm_out = nn.GroupNorm(num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6) + elif norm_layer == NormLayerType.PIXEL_NORM: + self.conv_norm_out = PixelNorm() + + self.conv_act = nn.SiLU() + self.conv_out = make_conv_nd( + dims=convolution_dimensions, + in_channels=feature_channels, + out_channels=out_channels, + kernel_size=3, + padding=1, + causal=True, + spatial_padding_mode=decoder_spatial_padding_mode, + ) + + if timestep_conditioning: + self.timestep_scale_multiplier = nn.Parameter(torch.tensor(1000.0)) + self.last_time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings( + embedding_dim=feature_channels * 2, size_emb_dim=0 + ) + self.last_scale_shift_table = nn.Parameter(torch.empty(2, feature_channels)) + + def forward( + self, + sample: torch.Tensor, + timestep: torch.Tensor | None = None, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + r""" + Decode latent representation into video frames. + Args: + sample: Latent tensor (B, 128, F', H', W'). + timestep: Timestep for conditioning (if timestep_conditioning=True). Uses default 0.05 if None. + generator: Random generator for deterministic noise injection (if inject_noise=True in blocks). + Returns: + Decoded video (B, 3, F, H, W) where F = 8x(F'-1) + 1, H = 32xH', W = 32xW'. + Example: (B, 128, 5, 16, 16) -> (B, 3, 33, 512, 512). + Note: First frame is removed after temporal upsampling regardless of causal mode. + When causal=False, allows future frame dependencies in convolutions but maintains same output shape. + """ + batch_size = sample.shape[0] + + # Add noise if timestep conditioning is enabled + if self.timestep_conditioning: + noise = ( + torch.randn( + sample.size(), + generator=generator, + dtype=sample.dtype, + device=sample.device, + ) + * self.decode_noise_scale + ) + + sample = noise + (1.0 - self.decode_noise_scale) * sample + + # Denormalize latents + sample = self.per_channel_statistics.un_normalize(sample) + + # Use default decode_timestep if timestep not provided + if timestep is None and self.timestep_conditioning: + timestep = torch.full((batch_size,), self.decode_timestep, device=sample.device, dtype=sample.dtype) + + sample = self.conv_in(sample, causal=self.causal) + + scaled_timestep = None + if self.timestep_conditioning: + if timestep is None: + raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True") + scaled_timestep = timestep * self.timestep_scale_multiplier.to(sample) + + for up_block in self.up_blocks: + if isinstance(up_block, UNetMidBlock3D): + block_kwargs = { + "causal": self.causal, + "timestep": scaled_timestep if self.timestep_conditioning else None, + "generator": generator, + } + sample = up_block(sample, **block_kwargs) + elif isinstance(up_block, ResnetBlock3D): + sample = up_block(sample, causal=self.causal, generator=generator) + else: + sample = up_block(sample, causal=self.causal) + + sample = self.conv_norm_out(sample) + + if self.timestep_conditioning: + embedded_timestep = self.last_time_embedder( + timestep=scaled_timestep.flatten(), + hidden_dtype=sample.dtype, + ) + embedded_timestep = embedded_timestep.view(batch_size, embedded_timestep.shape[-1], 1, 1, 1) + ada_values = self.last_scale_shift_table[None, ..., None, None, None].to( + device=sample.device, dtype=sample.dtype + ) + embedded_timestep.reshape( + batch_size, + 2, + -1, + embedded_timestep.shape[-3], + embedded_timestep.shape[-2], + embedded_timestep.shape[-1], + ) + shift, scale = ada_values.unbind(dim=1) + sample = sample * (1 + scale) + shift + + sample = self.conv_act(sample) + sample = self.conv_out(sample, causal=self.causal) + + # Final spatial expansion: reverse the initial patchify from encoder + # Moves pixels from channels back to spatial dimensions + # Example: (B, 48, F, 128, 128) -> (B, 3, F, 512, 512) with patch_size=4 + sample = unpatchify(sample, patch_size_hw=self.patch_size, patch_size_t=1) + + return sample + + def _prepare_tiles( + self, + latent: torch.Tensor, + tile_size: tuple[int, int] = (64, 512), + tile_stride: tuple[int, int] = (40, 448), + ) -> List[Tile]: + splitters = [DEFAULT_SPLIT_OPERATION] * len(latent.shape) + mappers = [DEFAULT_MAPPING_OPERATION] * len(latent.shape) + temporal_tile_size, spatial_tile_size = tile_size + temporal_tile_stride, spatial_tile_stride = tile_stride + spatial_overlap = spatial_tile_size - spatial_tile_stride + temporal_overlap = temporal_tile_size - temporal_tile_stride + + _validate_spatial_tiling(spatial_tile_size, spatial_overlap) + _validate_temporal_tiling(temporal_tile_size, temporal_overlap) + + if spatial_tile_size is not None: + long_side = max(latent.shape[3], latent.shape[4]) + + def enable_on_axis(axis_idx: int, factor: int) -> None: + size = spatial_tile_size // factor + overlap = spatial_overlap // factor + axis_length = latent.shape[axis_idx] + lower_threshold = max(2, overlap + 1) + tile_size = max(lower_threshold, round(size * axis_length / long_side)) + splitters[axis_idx] = split_with_symmetric_overlaps(tile_size, overlap) + mappers[axis_idx] = make_mapping_operation(map_spatial_interval_to_pixel, scale=factor) + + enable_on_axis(3, self.video_downscale_factors.height) + enable_on_axis(4, self.video_downscale_factors.width) + + if temporal_tile_size is not None: + tile_size_in_latents = temporal_tile_size // self.video_downscale_factors.time + overlap_in_latents = temporal_overlap // self.video_downscale_factors.time + splitters[2] = split_temporal_latents(tile_size_in_latents, overlap_in_latents) + mappers[2] = make_mapping_operation(map_temporal_interval_to_frame, scale=self.video_downscale_factors.time) + + return create_tiles(latent.shape, splitters, mappers) + + def tiled_decode( + self, + latent: torch.Tensor, + tile_size: tuple[int, int] = (64, 512), + tile_stride: tuple[int, int] = (40, 448), + timestep: torch.Tensor | None = None, + generator: torch.Generator | None = None, + ) -> Iterator[torch.Tensor]: + """ + Decode a latent tensor into video frames using tiled processing. + Splits the latent tensor into tiles, decodes each tile individually, + and yields video chunks as they become available. + Args: + latent: Input latent tensor (B, C, F', H', W'). + tile_size: Temporal/spatial tile size as `(frames, pixels)`. + tile_stride: Temporal/spatial tile stride as `(frames, pixels)`. + timestep: Optional timestep for decoder conditioning. + generator: Optional random generator for deterministic decoding. + Yields: + Video chunks (B, C, T, H, W) by temporal slices; + """ + + # Calculate full video shape from latent shape to get spatial dimensions + full_video_shape = VideoLatentShape.from_torch_shape(latent.shape).upscale(self.video_downscale_factors) + tiles = self._prepare_tiles(latent, tile_size=tile_size, tile_stride=tile_stride) + + temporal_groups = self._group_tiles_by_temporal_slice(tiles) + + # State for temporal overlap handling + previous_chunk = None + previous_weights = None + previous_temporal_slice = None + + for temporal_group_tiles in temporal_groups: + curr_temporal_slice = temporal_group_tiles[0].out_coords[2] + + # Calculate the shape of the temporal buffer for this group of tiles. + # The temporal length depends on whether this is the first tile (starts at 0) or not. + # - First tile: (frames - 1) * scale + 1 + # - Subsequent tiles: frames * scale + # This logic is handled by TemporalAxisMapping and reflected in out_coords. + temporal_tile_buffer_shape = full_video_shape._replace( + frames=curr_temporal_slice.stop - curr_temporal_slice.start, + ) + + buffer = torch.zeros( + temporal_tile_buffer_shape.to_torch_shape(), + device=latent.device, + dtype=latent.dtype, + ) + + curr_weights = self._accumulate_temporal_group_into_buffer( + group_tiles=temporal_group_tiles, + buffer=buffer, + latent=latent, + timestep=timestep, + generator=generator, + ) + + # Blend with previous temporal chunk if it exists + if previous_chunk is not None: + # Check if current temporal slice overlaps with previous temporal slice + if previous_temporal_slice.stop > curr_temporal_slice.start: + overlap_len = previous_temporal_slice.stop - curr_temporal_slice.start + temporal_overlap_slice = slice(curr_temporal_slice.start - previous_temporal_slice.start, None) + + # The overlap is already masked before it reaches this step. Each tile is accumulated into buffer + # with its trapezoidal mask, and curr_weights accumulates the same mask. In the overlap blend we add + # the masked values (buffer[...]) and the corresponding weights (curr_weights[...]) into the + # previous buffers, then later normalize by weights. + previous_chunk[:, :, temporal_overlap_slice, :, :] += buffer[:, :, slice(0, overlap_len), :, :] + previous_weights[:, :, temporal_overlap_slice, :, :] += curr_weights[ + :, :, slice(0, overlap_len), :, : + ] + + buffer[:, :, slice(0, overlap_len), :, :] = previous_chunk[:, :, temporal_overlap_slice, :, :] + curr_weights[:, :, slice(0, overlap_len), :, :] = previous_weights[ + :, :, temporal_overlap_slice, :, : + ] + + # Yield the non-overlapping part of the previous chunk + previous_weights = previous_weights.clamp(min=1e-8) + yield_len = curr_temporal_slice.start - previous_temporal_slice.start + yield (previous_chunk / previous_weights)[:, :, :yield_len, :, :] + + # Update state for next iteration + previous_chunk = buffer + previous_weights = curr_weights + previous_temporal_slice = curr_temporal_slice + + # Yield any remaining chunk + if previous_chunk is not None: + previous_weights = previous_weights.clamp(min=1e-8) + yield previous_chunk / previous_weights + + def _group_tiles_by_temporal_slice(self, tiles: List[Tile]) -> List[List[Tile]]: + """Group tiles by their temporal output slice.""" + if not tiles: + return [] + + groups = [] + current_slice = tiles[0].out_coords[2] + current_group = [] + + for tile in tiles: + tile_slice = tile.out_coords[2] + if tile_slice == current_slice: + current_group.append(tile) + else: + groups.append(current_group) + current_slice = tile_slice + current_group = [tile] + + # Add the final group + if current_group: + groups.append(current_group) + + return groups + + def _accumulate_temporal_group_into_buffer( + self, + group_tiles: List[Tile], + buffer: torch.Tensor, + latent: torch.Tensor, + timestep: torch.Tensor | None, + generator: torch.Generator | None, + ) -> torch.Tensor: + """ + Decode and accumulate all tiles of a temporal group into a local buffer. + The buffer is local to the group and always starts at time 0; temporal coordinates + are rebased by subtracting temporal_slice.start. + """ + temporal_slice = group_tiles[0].out_coords[2] + + weights = torch.zeros_like(buffer) + + for tile in group_tiles: + decoded_tile = self.forward(latent[tile.in_coords], timestep, generator) + mask = tile.blend_mask.to(device=buffer.device, dtype=buffer.dtype) + temporal_offset = tile.out_coords[2].start - temporal_slice.start + # Use the tile's output coordinate length, not the decoded tile's length, + # as the decoder may produce a different number of frames than expected + expected_temporal_len = tile.out_coords[2].stop - tile.out_coords[2].start + decoded_temporal_len = decoded_tile.shape[2] + + # Ensure we don't exceed the buffer or decoded tile bounds + actual_temporal_len = min(expected_temporal_len, decoded_temporal_len, buffer.shape[2] - temporal_offset) + + chunk_coords = ( + slice(None), # batch + slice(None), # channels + slice(temporal_offset, temporal_offset + actual_temporal_len), + tile.out_coords[3], # height + tile.out_coords[4], # width + ) + + # Slice decoded_tile and mask to match the actual length we're writing + decoded_slice = decoded_tile[:, :, :actual_temporal_len, :, :] + mask_slice = mask[:, :, :actual_temporal_len, :, :] if mask.shape[2] > 1 else mask + + buffer[chunk_coords] += decoded_slice * mask_slice + weights[chunk_coords] += mask_slice + + return weights + + +def decode_video( + latent: torch.Tensor, + video_decoder: VideoDecoder, + tiled: bool = False, + tile_size: tuple[int, int] = (64, 512), + tile_stride: tuple[int, int] = (40, 448), + generator: torch.Generator | None = None, +) -> Iterator[torch.Tensor]: + """ + Decode a video latent tensor with the given decoder. + Args: + latent: Tensor [c, f, h, w] + video_decoder: Decoder module. + tiled: Whether to enable tiled decoding. + tile_size: Temporal/spatial tile size as `(frames, pixels)`. + tile_stride: Temporal/spatial tile stride as `(frames, pixels)`. + generator: Optional random generator for deterministic decoding. + Yields: + Decoded chunk [b, c, f, h, w] in the decoder output range. + """ + + if tiled: + for frames in video_decoder.tiled_decode( + latent, tile_size=tile_size, tile_stride=tile_stride, generator=generator + ): + yield frames + else: + decoded_video = video_decoder(latent, generator=generator) + yield decoded_video + + +def get_video_chunks_number( + num_frames: int, + tiled: bool = False, + tile_size: tuple[int, int] = (64, 512), + tile_stride: tuple[int, int] = (40, 448), +) -> int: + """ + Get the number of video chunks for a given number of frames and tiling configuration. + Args: + num_frames: Number of frames in the video. + tiled: Whether tiled decoding is enabled. + tile_size: Temporal/spatial tile size as `(frames, pixels)`. + tile_stride: Temporal/spatial tile stride as `(frames, pixels)`. + Returns: + Number of video chunks. + """ + if not tiled: + return 1 + frame_stride = tile_stride[0] + return (num_frames - 1 + frame_stride - 1) // frame_stride + + +def split_with_symmetric_overlaps(size: int, overlap: int) -> SplitOperation: + def split(dimension_size: int) -> DimensionIntervals: + if dimension_size <= size: + return DEFAULT_SPLIT_OPERATION(dimension_size) + amount = (dimension_size + size - 2 * overlap - 1) // (size - overlap) + starts = [i * (size - overlap) for i in range(amount)] + ends = [start + size for start in starts] + ends[-1] = dimension_size + left_ramps = [0] + [overlap] * (amount - 1) + right_ramps = [overlap] * (amount - 1) + [0] + return DimensionIntervals(starts=starts, ends=ends, left_ramps=left_ramps, right_ramps=right_ramps) + + return split + + +def split_temporal_latents(size: int, overlap: int) -> SplitOperation: + """Split a temporal axis into overlapping tiles with causal handling. + Example with size=24, overlap=8 (units are whatever axis you split): + Non-causal split would produce: + Tile 0: [0, 24), left_ramp=0, right_ramp=8 + Tile 1: [16, 40), left_ramp=8, right_ramp=8 + Tile 2: [32, 56), left_ramp=8, right_ramp=0 + Causal split produces: + Tile 0: [0, 24), left_ramp=0, right_ramp=8 (unchanged - starts at anchor) + Tile 1: [15, 40), left_ramp=9, right_ramp=8 (shifted back 1, ramp +1) + Tile 2: [31, 56), left_ramp=9, right_ramp=0 (shifted back 1, ramp +1) + This ensures each tile can causally depend on frames from previous tiles while maintaining + proper temporal continuity through the blend ramps. + Args: + size: Tile size in *axis units* (latent steps for LTX time tiling) + overlap: Overlap between tiles in the same units + Returns: + Split operation that divides temporal dimension with causal handling + """ + non_causal_split = split_with_symmetric_overlaps(size, overlap) + + def split(dimension_size: int) -> DimensionIntervals: + if dimension_size <= size: + return DEFAULT_SPLIT_OPERATION(dimension_size) + intervals = non_causal_split(dimension_size) + + starts = intervals.starts + starts[1:] = [s - 1 for s in starts[1:]] + + # Extend blend ramps by 1 for non-first tiles to blend over the extra frame + left_ramps = intervals.left_ramps + left_ramps[1:] = [r + 1 for r in left_ramps[1:]] + + return replace(intervals, starts=starts, left_ramps=left_ramps) + + return split + + +def split_temporal_frames(tile_size_frames: int, overlap_frames: int) -> SplitOperation: + """Split a temporal axis in video frame space into overlapping tiles. + Args: + tile_size_frames: Tile length in frames. + overlap_frames: Overlap between consecutive tiles in frames. + Returns: + Split operation that takes frame count and returns DimensionIntervals in frame indices. + """ + non_causal_split = split_with_symmetric_overlaps(tile_size_frames, overlap_frames) + + def split(dimension_size: int) -> DimensionIntervals: + if dimension_size <= tile_size_frames: + return DEFAULT_SPLIT_OPERATION(dimension_size) + intervals = non_causal_split(dimension_size) + ends = intervals.ends + ends[:-1] = [e + 1 for e in ends[:-1]] + right_ramps = [0] * len(intervals.right_ramps) + return replace(intervals, ends=ends, right_ramps=right_ramps) + + return split + + +def make_mapping_operation( + map_func: Callable[[int, int, int, int, int], Tuple[slice, torch.Tensor | None]], + scale: int, +) -> MappingOperation: + """Create a mapping operation over a set of tiling intervals. + The given mapping function is applied to each interval in the input dimension. The result function is used for + creating tiles in the output dimension. + Args: + map_func: Mapping function to create the mapping operation from + scale: Scale factor for the transformation, used as an argument for the mapping function + Returns: + Mapping operation that takes a set of tiling intervals and returns a set of slices and masks in the output + dimension. + """ + + def map_op(intervals: DimensionIntervals) -> tuple[list[slice], list[torch.Tensor | None]]: + output_slices: list[slice] = [] + masks_1d: list[torch.Tensor | None] = [] + number_of_slices = len(intervals.starts) + for i in range(number_of_slices): + start = intervals.starts[i] + end = intervals.ends[i] + left_ramp = intervals.left_ramps[i] + right_ramp = intervals.right_ramps[i] + output_slice, mask_1d = map_func(start, end, left_ramp, right_ramp, scale) + output_slices.append(output_slice) + masks_1d.append(mask_1d) + return output_slices, masks_1d + + return map_op + + +def map_temporal_interval_to_frame( + begin: int, + end: int, + left_ramp: int, + right_ramp: int, + scale: int, +) -> Tuple[slice, torch.Tensor]: + """Map temporal interval in latent space to video frame space. + Args: + begin: Start position in latent space + end: End position in latent space + left_ramp: Left ramp size in latent space + right_ramp: Right ramp size in latent space + scale: Scale factor for transformation + Returns: + Tuple of (output_slice, blend_mask) + """ + start = begin * scale + stop = 1 + (end - 1) * scale + + left_ramp_frames = 0 if left_ramp == 0 else 1 + (left_ramp - 1) * scale + right_ramp_frames = right_ramp * scale + + mask_1d = compute_trapezoidal_mask_1d(stop - start, left_ramp_frames, right_ramp_frames, True) + return slice(start, stop), mask_1d + + +def map_temporal_interval_to_latent( + begin: int, end: int, left_ramp: int, right_ramp: int | None = None, scale: int = 1 +) -> Tuple[slice, torch.Tensor]: + """ + Map temporal interval in video frame space to latent space. + Args: + begin: Start position in video frame space + end: End position in video frame space + left_ramp: Left ramp size in video frame space + right_ramp: Right ramp size in video frame space + scale: Scale factor for transformation + Returns: + Tuple of (output_slice, blend_mask) + """ + start = begin // scale + stop = (end - 1) // scale + 1 + + left_ramp_latents = 0 if left_ramp == 0 else 1 + (left_ramp - 1) // scale + right_ramp_latents = right_ramp // scale + + if right_ramp_latents != 0: + raise ValueError("For tiled encoding, temporal tiles are expected to have a right ramp equal to 0") + + mask_1d = compute_rectangular_mask_1d(stop - start, left_ramp_latents, right_ramp_latents) + + return slice(start, stop), mask_1d + + +def map_spatial_interval_to_pixel( + begin: int, + end: int, + left_ramp: int, + right_ramp: int, + scale: int, +) -> Tuple[slice, torch.Tensor]: + """Map spatial interval in latent space to pixel space. + Args: + begin: Start position in latent space + end: End position in latent space + left_ramp: Left ramp size in latent space + right_ramp: Right ramp size in latent space + scale: Scale factor for transformation + """ + start = begin * scale + stop = end * scale + mask_1d = compute_trapezoidal_mask_1d(stop - start, left_ramp * scale, right_ramp * scale, False) + return slice(start, stop), mask_1d + + +def map_spatial_interval_to_latent( + begin: int, + end: int, + left_ramp: int, + right_ramp: int, + scale: int, +) -> Tuple[slice, torch.Tensor]: + """Map spatial interval in pixel space to latent space. + Args: + begin: Start position in pixel space + end: End position in pixel space + left_ramp: Left ramp size in pixel space + right_ramp: Right ramp size in pixel space + scale: Scale factor for transformation + Returns: + Tuple of (output_slice, blend_mask) + """ + start = begin // scale + stop = end // scale + left_ramp = max(0, left_ramp // scale - 1) + + right_ramp = 0 if right_ramp == 0 else 1 + + mask_1d = compute_rectangular_mask_1d(stop - start, left_ramp, right_ramp) + return slice(start, stop), mask_1d + + +class VideoEncoderConfigurator: + """Thin config adapter for constructing `VideoEncoder`.""" + + @classmethod + def from_config(cls: type[VideoEncoder], config: dict) -> VideoEncoder: + return VideoEncoder(**_ltx_video_vae_encoder_kwargs_from_config(config)) + + +class VideoDecoderConfigurator: + """Thin config adapter for constructing `VideoDecoder`.""" + + @classmethod + def from_config(cls: type[VideoDecoder], config: dict) -> VideoDecoder: + return VideoDecoder(**_ltx_video_vae_decoder_kwargs_from_config(config)) + + +def _ltx_video_vae_config_section(config: dict) -> dict: + return config.get("vae", config) + + +def _ltx_video_vae_encoder_kwargs_from_config(config: dict) -> dict: + config = _ltx_video_vae_config_section(config) + return { + "convolution_dimensions": config.get("dims", 3), + "in_channels": config.get("in_channels", 3), + "out_channels": config.get("latent_channels", 128), + "encoder_blocks": config.get("encoder_blocks", []), + "patch_size": config.get("patch_size", 4), + "norm_layer": NormLayerType(config.get("norm_layer", "pixel_norm")), + "latent_log_var": LogVarianceType(config.get("latent_log_var", "uniform")), + "encoder_spatial_padding_mode": PaddingModeType(config.get("spatial_padding_mode", "zeros")), + } + + +def _ltx_video_vae_decoder_kwargs_from_config(config: dict) -> dict: + config = _ltx_video_vae_config_section(config) + return { + "convolution_dimensions": config.get("dims", 3), + "in_channels": config.get("latent_channels", 128), + "out_channels": config.get("out_channels", 3), + "decoder_blocks": config.get("decoder_blocks", []), + "patch_size": config.get("patch_size", 4), + "norm_layer": NormLayerType(config.get("norm_layer", "pixel_norm")), + "causal": config.get("causal_decoder", False), + "timestep_conditioning": config.get("timestep_conditioning", False), + "decoder_spatial_padding_mode": PaddingModeType(config.get("spatial_padding_mode", "zeros")), + "base_channels": config.get("decoder_base_channels", 128), + } + + +def _ltx_video_vae_kwargs_from_config(config: dict) -> dict: + config = _ltx_video_vae_config_section(config) + return { + "convolution_dimensions": config.get("dims", 3), + "in_channels": config.get("in_channels", 3), + "out_channels": config.get("out_channels", 3), + "latent_channels": config.get("latent_channels", 128), + "encoder_blocks": tuple(config.get("encoder_blocks", DEFAULT_ENCODER_BLOCKS)), + "decoder_blocks": tuple(config.get("decoder_blocks", DEFAULT_DECODER_BLOCKS)), + "patch_size": config.get("patch_size", 4), + "norm_layer": NormLayerType(config.get("norm_layer", "pixel_norm")), + "latent_log_var": LogVarianceType(config.get("latent_log_var", "uniform")), + "causal_decoder": config.get("causal_decoder", False), + "timestep_conditioning": config.get("timestep_conditioning", False), + "spatial_padding_mode": PaddingModeType(config.get("spatial_padding_mode", "zeros")), + "decoder_base_channels": config.get("decoder_base_channels", 128), + } + + +class LTXVideoVAE(BaseModel): + """Registered LTX 2.3 VAE wrapper with shared encoder/decoder checkpoint loading.""" + + def __init__( + self, + convolution_dimensions: int = 3, + in_channels: int = 3, + out_channels: int = 3, + latent_channels: int = 128, + encoder_blocks: tuple[tuple[str, int | dict], ...] = DEFAULT_ENCODER_BLOCKS, + decoder_blocks: tuple[tuple[str, int | dict], ...] = DEFAULT_DECODER_BLOCKS, + patch_size: int = 4, + norm_layer: NormLayerType = NormLayerType.PIXEL_NORM, + latent_log_var: LogVarianceType = LogVarianceType.UNIFORM, + causal_decoder: bool = False, + timestep_conditioning: bool = False, + spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, + decoder_base_channels: int = 128, + ) -> None: + super().__init__() + self.encoder = VideoEncoder( + convolution_dimensions=convolution_dimensions, + in_channels=in_channels, + out_channels=latent_channels, + encoder_blocks=list(encoder_blocks), + patch_size=patch_size, + norm_layer=norm_layer, + latent_log_var=latent_log_var, + encoder_spatial_padding_mode=spatial_padding_mode, + ) + self.decoder = VideoDecoder( + convolution_dimensions=convolution_dimensions, + in_channels=latent_channels, + out_channels=out_channels, + decoder_blocks=list(decoder_blocks), + patch_size=patch_size, + norm_layer=norm_layer, + causal=causal_decoder, + timestep_conditioning=timestep_conditioning, + decoder_spatial_padding_mode=spatial_padding_mode, + base_channels=decoder_base_channels, + ) + + @classmethod + def from_config(cls, config: dict) -> LTXVideoVAE: + """Build the registered LTX VAE from a config dictionary.""" + return cls(**_ltx_video_vae_kwargs_from_config(config)) + + @property + def per_channel_statistics(self) -> nn.Module: + return self.encoder.per_channel_statistics + + def set_parallelism(self, parallelism: int): + self.parallelism = parallelism + + def encode( + self, + video: torch.Tensor, + tiled: bool = False, + tile_size: tuple[int, int] = (64, 512), + tile_stride: tuple[int, int] = (40, 448), + ) -> torch.Tensor: + if not tiled: + return self.encoder(video) + return self.encoder.tiled_encode(video, tile_size=tile_size, tile_stride=tile_stride) + + def decode( + self, + latent: torch.Tensor, + tiled: bool = False, + tile_size: tuple[int, int] = (64, 512), + tile_stride: tuple[int, int] = (40, 448), + timestep: torch.Tensor | None = None, + generator: torch.Generator | None = None, + ) -> Iterator[torch.Tensor]: + if not tiled: + yield self.decoder(latent, timestep=timestep, generator=generator) + return + yield from self.decoder.tiled_decode( + latent, + tile_size=tile_size, + tile_stride=tile_stride, + timestep=timestep, + generator=generator, + ) + + @staticmethod + def state_dict_converter(): + return LTXVideoVAEStateDictConverter() + + def enable_sequential_cpu_offload( + self, + device: torch.device | None = None, + torch_dtype: torch.dtype | None = None, + ) -> None: + """Enable sequential CPU offload for the encoder and decoder stacks.""" + from telefuser.offload import ( + AutoWrappedLinear, + AutoWrappedModule, + enable_sequential_cpu_offload, + ) + + if device is None: + device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + if torch_dtype is None: + torch_dtype = next(iter(self.parameters())).dtype + + dtype = next(iter(self.parameters())).dtype + enable_sequential_cpu_offload( + self, + module_map={ + torch.nn.Linear: AutoWrappedLinear, + torch.nn.Conv2d: AutoWrappedModule, + torch.nn.Conv3d: AutoWrappedModule, + torch.nn.GroupNorm: AutoWrappedModule, + PixelNorm: AutoWrappedModule, + }, + module_config=dict( + offload_dtype=dtype, + offload_device="cpu", + onload_dtype=dtype, + onload_device=device, + computation_dtype=torch_dtype, + computation_device=device, + ), + ) + + +class LTXVideoVAEStateDictConverter: + """Convert the monolithic LTX checkpoint into the registered VAE wrapper layout.""" + + _ENCODER_PREFIX = "vae.encoder." + _DECODER_PREFIX = "vae.decoder." + _STATS_PREFIX = "vae.per_channel_statistics." + + def from_official(self, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + converted_state_dict: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if key.startswith(self._ENCODER_PREFIX): + converted_state_dict[f"encoder.{key.removeprefix(self._ENCODER_PREFIX)}"] = value + elif key.startswith(self._DECODER_PREFIX): + converted_state_dict[f"decoder.{key.removeprefix(self._DECODER_PREFIX)}"] = value + elif key.startswith(self._STATS_PREFIX): + suffix = key.removeprefix(self._STATS_PREFIX) + converted_state_dict[f"encoder.per_channel_statistics.{suffix}"] = value + converted_state_dict[f"decoder.per_channel_statistics.{suffix}"] = value + return converted_state_dict + + def from_diffusers(self, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + raise NotImplementedError("LTXVideoVAE only supports the civitai-style single-file checkpoint.") + + +LTXVideoDecoder = VideoDecoder +LTXVideoEncoder = VideoEncoder + + +class LTX25ConvVideoVAE(LTXVideoVAE): + """Isolated LTX-2.5 Conv VAE wrapper with strict split-checkpoint loading.""" + + @classmethod + def from_config(cls, config: dict) -> "LTX25ConvVideoVAE": + return cls(**_ltx_video_vae_kwargs_from_config(config)) + + @classmethod + def from_checkpoint( + cls, + checkpoint_path: str | Path, + *, + device: torch.device | str = "cpu", + torch_dtype: torch.dtype = torch.bfloat16, + ) -> "LTX25ConvVideoVAE": + """Construct and strictly load the official standalone Conv VAE checkpoint.""" + checkpoint = inspect_checkpoint(checkpoint_path) + with torch.device("meta"): + model = cls.from_config(checkpoint.config) + unexpected, missing = ltx25_conv_video_vae_checkpoint_key_coverage(checkpoint.path, set(model.state_dict())) + if unexpected or missing: + raise ValueError( + "LTX-2.5 Conv VAE checkpoint coverage mismatch: " + f"unexpected={sorted(unexpected)[:5]}, missing={sorted(missing)[:5]}" + ) + state_dict = _load_ltx25_conv_video_vae_state_dict(checkpoint.path) + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=True, assign=True) + if missing_keys or unexpected_keys: + raise ValueError( + f"LTX-2.5 Conv VAE load mismatch: missing={missing_keys[:5]}, unexpected={unexpected_keys[:5]}" + ) + model = model.to(device=device, dtype=torch_dtype).eval() + from .memory_efficient_decode import ( + convert_decoder_weights_to_channels_last_3d, + enable_memory_efficient_decode, + ) + + convert_decoder_weights_to_channels_last_3d(model.decoder) + enable_memory_efficient_decode(model.decoder) + return model + + +def _conv_video_vae_key_to_model_keys(key: str) -> Iterable[str]: + if key.startswith("encoder.") or key.startswith("decoder."): + return (key,) + if key.startswith("per_channel_statistics."): + suffix = key.removeprefix("per_channel_statistics.") + return (f"encoder.per_channel_statistics.{suffix}", f"decoder.per_channel_statistics.{suffix}") + return () + + +def _load_ltx25_conv_video_vae_state_dict(checkpoint_path: str | Path) -> dict[str, torch.Tensor]: + state_dict: dict[str, torch.Tensor] = {} + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as checkpoint: + for key in checkpoint.keys(): + targets = tuple(_conv_video_vae_key_to_model_keys(key)) + if not targets: + continue + value = checkpoint.get_tensor(key) + for target in targets: + state_dict[target] = value + return state_dict + + +def ltx25_conv_video_vae_checkpoint_key_coverage( + checkpoint_path: str | Path, + model_keys: set[str], +) -> tuple[set[str], set[str]]: + """Return unexplained checkpoint keys and missing Conv VAE model keys.""" + mapped: set[str] = set() + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as checkpoint: + for key in checkpoint.keys(): + mapped.update(_conv_video_vae_key_to_model_keys(key)) + return mapped - model_keys, model_keys - mapped + + +__all__ = [ + "LTXVideoVAE", + "LTX25ConvVideoVAE", + "LTXVideoVAEStateDictConverter", + "LTXVideoDecoder", + "LTXVideoEncoder", + "LogVarianceType", + "NormLayerType", + "PaddingModeType", + "PixelNorm", + "SpatioTemporalScaleFactors", + "VIDEO_SCALE_FACTORS", + "VideoDecoder", + "VideoDecoderConfigurator", + "VideoEncoder", + "VideoEncoderConfigurator", + "VideoLatentShape", + "VideoPixelShape", + "build_normalization_layer", + "decode_video", + "get_video_chunks_number", + "ltx25_conv_video_vae_checkpoint_key_coverage", +] diff --git a/telefuser/models/ltx25/diff_vae/__init__.py b/telefuser/models/ltx25/diff_vae/__init__.py new file mode 100644 index 0000000..20baf58 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/__init__.py @@ -0,0 +1,5 @@ +"""Isolated eager LTX-2.5 diffusion video VAE decoder.""" + +from .diffusion_video_decoder import DiffusionVideoDecoder + +__all__ = ["DiffusionVideoDecoder"] diff --git a/telefuser/models/ltx25/diff_vae/convolution.py b/telefuser/models/ltx25/diff_vae/convolution.py new file mode 100644 index 0000000..fbb8a05 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/convolution.py @@ -0,0 +1,317 @@ +from typing import Tuple, Union + +import torch +from einops import rearrange +from torch import nn +from torch.nn import functional as F + +from telefuser.models.ltx25.diff_vae.enums import PaddingModeType + + +def make_conv_nd( # noqa: PLR0913 + dims: Union[int, Tuple[int, int]], + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + padding: int = 0, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + causal: bool = False, + spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, + temporal_padding_mode: PaddingModeType = PaddingModeType.ZEROS, +) -> nn.Module: + if not (spatial_padding_mode == temporal_padding_mode or causal): + raise NotImplementedError("spatial and temporal padding modes must be equal") + if dims == 2: + return nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + bias=bias, + padding_mode=spatial_padding_mode.value, + ) + elif dims == 3: + if causal: + return CausalConv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + groups=groups, + bias=bias, + spatial_padding_mode=spatial_padding_mode, + ) + return nn.Conv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + bias=bias, + padding_mode=spatial_padding_mode.value, + ) + elif dims == (2, 1): + return DualConv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + bias=bias, + padding_mode=spatial_padding_mode.value, + ) + else: + raise ValueError(f"unsupported dimensions: {dims}") + + +def make_linear_nd( + dims: int, + in_channels: int, + out_channels: int, + bias: bool = True, +) -> nn.Module: + if dims == 2: + return nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias) + elif dims in (3, (2, 1)): + return nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias) + else: + raise ValueError(f"unsupported dimensions: {dims}") + + +class DualConv3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: Union[int, Tuple[int, int, int]] = 1, + padding: Union[int, Tuple[int, int, int]] = 0, + dilation: Union[int, Tuple[int, int, int]] = 1, + groups: int = 1, + bias: bool = True, + padding_mode: str = "zeros", + ) -> None: + super(DualConv3d, self).__init__() + + self.in_channels = in_channels + self.out_channels = out_channels + self.padding_mode = padding_mode + # Ensure kernel_size, stride, padding, and dilation are tuples of length 3 + if isinstance(kernel_size, int): + kernel_size = (kernel_size, kernel_size, kernel_size) + if kernel_size == (1, 1, 1): + raise ValueError("kernel_size must be greater than 1. Use make_linear_nd instead.") + if isinstance(stride, int): + stride = (stride, stride, stride) + if isinstance(padding, int): + padding = (padding, padding, padding) + if isinstance(dilation, int): + dilation = (dilation, dilation, dilation) + + # Set parameters for convolutions + self.groups = groups + self.bias = bias + + # Define the size of the channels after the first convolution + intermediate_channels = out_channels if in_channels < out_channels else in_channels + + # Define parameters for the first convolution + self.weight1 = nn.Parameter( + torch.Tensor( + intermediate_channels, + in_channels // groups, + 1, + kernel_size[1], + kernel_size[2], + ) + ) + self.stride1 = (1, stride[1], stride[2]) + self.padding1 = (0, padding[1], padding[2]) + self.dilation1 = (1, dilation[1], dilation[2]) + if bias: + self.bias1 = nn.Parameter(torch.Tensor(intermediate_channels)) + else: + self.register_parameter("bias1", None) + + # Define parameters for the second convolution + self.weight2 = nn.Parameter(torch.Tensor(out_channels, intermediate_channels // groups, kernel_size[0], 1, 1)) + self.stride2 = (stride[0], 1, 1) + self.padding2 = (padding[0], 0, 0) + self.dilation2 = (dilation[0], 1, 1) + if bias: + self.bias2 = nn.Parameter(torch.Tensor(out_channels)) + else: + self.register_parameter("bias2", None) + + # Initialize weights and biases + self.reset_parameters() + + def reset_parameters(self) -> None: + nn.init.kaiming_uniform_(self.weight1, a=torch.sqrt(5)) + nn.init.kaiming_uniform_(self.weight2, a=torch.sqrt(5)) + if self.bias: + fan_in1, _ = nn.init._calculate_fan_in_and_fan_out(self.weight1) + bound1 = 1 / torch.sqrt(fan_in1) + nn.init.uniform_(self.bias1, -bound1, bound1) + fan_in2, _ = nn.init._calculate_fan_in_and_fan_out(self.weight2) + bound2 = 1 / torch.sqrt(fan_in2) + nn.init.uniform_(self.bias2, -bound2, bound2) + + def forward( + self, + x: torch.Tensor, + use_conv3d: bool = False, + skip_time_conv: bool = False, + ) -> torch.Tensor: + if use_conv3d: + return self.forward_with_3d(x=x, skip_time_conv=skip_time_conv) + else: + return self.forward_with_2d(x=x, skip_time_conv=skip_time_conv) + + def forward_with_3d(self, x: torch.Tensor, skip_time_conv: bool = False) -> torch.Tensor: + # First convolution + x = F.conv3d( + x, + self.weight1, + self.bias1, + self.stride1, + self.padding1, + self.dilation1, + self.groups, + padding_mode=self.padding_mode, + ) + + if skip_time_conv: + return x + + # Second convolution + x = F.conv3d( + x, + self.weight2, + self.bias2, + self.stride2, + self.padding2, + self.dilation2, + self.groups, + padding_mode=self.padding_mode, + ) + + return x + + def forward_with_2d(self, x: torch.Tensor, skip_time_conv: bool = False) -> torch.Tensor: + b, _, _, h, w = x.shape + + # First 2D convolution + x = rearrange(x, "b c d h w -> (b d) c h w") + # Squeeze the depth dimension out of weight1 since it's 1 + weight1 = self.weight1.squeeze(2) + # Select stride, padding, and dilation for the 2D convolution + stride1 = (self.stride1[1], self.stride1[2]) + padding1 = (self.padding1[1], self.padding1[2]) + dilation1 = (self.dilation1[1], self.dilation1[2]) + x = F.conv2d( + x, + weight1, + self.bias1, + stride1, + padding1, + dilation1, + self.groups, + padding_mode=self.padding_mode, + ) + + _, _, h, w = x.shape + + if skip_time_conv: + x = rearrange(x, "(b d) c h w -> b c d h w", b=b) + return x + + # Second convolution which is essentially treated as a 1D convolution across the 'd' dimension + x = rearrange(x, "(b d) c h w -> (b h w) c d", b=b) + + # Reshape weight2 to match the expected dimensions for conv1d + weight2 = self.weight2.squeeze(-1).squeeze(-1) + # Use only the relevant dimension for stride, padding, and dilation for the 1D convolution + stride2 = self.stride2[0] + padding2 = self.padding2[0] + dilation2 = self.dilation2[0] + x = F.conv1d( + x, + weight2, + self.bias2, + stride2, + padding2, + dilation2, + self.groups, + padding_mode=self.padding_mode, + ) + x = rearrange(x, "(b h w) c d -> b c d h w", b=b, h=h, w=w) + + return x + + @property + def weight(self) -> torch.Tensor: + return self.weight2 + + +class CausalConv3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int = 3, + stride: Union[int, Tuple[int]] = 1, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, + ) -> None: + super().__init__() + + self.in_channels = in_channels + self.out_channels = out_channels + + kernel_size = (kernel_size, kernel_size, kernel_size) + self.time_kernel_size = kernel_size[0] + + dilation = (dilation, 1, 1) + + height_pad = kernel_size[1] // 2 + width_pad = kernel_size[2] // 2 + padding = (0, height_pad, width_pad) + + self.conv = nn.Conv3d( + in_channels, + out_channels, + kernel_size, + stride=stride, + dilation=dilation, + padding=padding, + padding_mode=spatial_padding_mode.value, + groups=groups, + bias=bias, + ) + + def forward(self, x: torch.Tensor, causal: bool = True) -> torch.Tensor: + if causal: + first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, self.time_kernel_size - 1, 1, 1)) + x = torch.concatenate((first_frame_pad, x), dim=2) + else: + first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1)) + last_frame_pad = x[:, :, -1:, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1)) + x = torch.concatenate((first_frame_pad, x, last_frame_pad), dim=2) + x = self.conv(x) + return x + + @property + def weight(self) -> torch.Tensor: + return self.conv.weight diff --git a/telefuser/models/ltx25/diff_vae/diffusion_tiling.py b/telefuser/models/ltx25/diff_vae/diffusion_tiling.py new file mode 100644 index 0000000..3b300af --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/diffusion_tiling.py @@ -0,0 +1,771 @@ +"""DiffVAE tiling helpers: schedule, pad/crop/size-floor, blend utilities. +Decode orchestration lives on ``DiffusionVideoDecoder``. This module owns the +geometry/schedule/mask pieces that tiling uses. +""" + +from __future__ import annotations + +import itertools +import math +from dataclasses import dataclass +from typing import List, Literal, Sequence, Tuple + +import torch + +from telefuser.models.ltx25.diff_vae.transformer.config import DiffVAEMode + +from .tiling import ( + DEFAULT_SPLIT_OPERATION, + DimensionInterval, + DimensionSizeConfig, + SplitOperation, + Tile, + TileSizeConfig, + TilingConfig, + _validate_overlap, + compute_trapezoidal_mask_1d, + split_by_size, + untiled_mask_1d, +) +from .types import VIDEO_SCALE_FACTORS, SpatioTemporalScaleFactors + +ResizeAxisMode = Literal["repeat_last", "symmetric"] + +# Peak-activation heuristic (bytes): +# hard (resident for whole tiled decode): +# stage-4 input feature (stages 1-3 output): s4_txs4_hxs4_wxstage4_channelsxbf16 +# per temporal group + stage-5 tile: +# accumulator: H x W x (2 x tile_t) x out_channels x accum_elem +# (full-frame spatial; ~current group buffer + retained overlap stub / still- +# live emitted exclusive chunk heading to encode; RGB by default; accum_elem +# matches decode: fp16 when features are bf16, else feature dtype) +# stage-5: stage5_tokens x stage5_channels x bf16 x coef +# where stage5_tokens = F x (H/patch) x (W/patch) and coef folds NA working-set multiplicity. +_MEM_COEF_BY_MODE: dict[DiffVAEMode, float] = { + DiffVAEMode.COMBINED_COMPILE: 11, + DiffVAEMode.CHUNKED_COMPILE: 7, # natten backend; Triton/eager use CHUNKED_EAGER + DiffVAEMode.CHUNKED_EAGER: 5, + DiffVAEMode.BLACKWELL_DSL: 2.5, +} +_DEFAULT_ELEMENT_SIZE: int = 2 # bf16 features → fp16 accumulator / bf16 stage-5 +_ACCUMULATOR_CHANNELS: int = 3 # RGB pixel blend buffer (decoder out_channels) +_MIN_MODEL_BYTES_FLOOR: int = 1 << 30 # never assume a free DiffVAE weight footprint +_BUDGET_SAFETY_BYTES_EAGER: int = 1 << 30 +_BUDGET_SAFETY_BYTES_COMPILED: int = 2 << 30 +# Empirical H100 guardrail for the upstream CHUNKED_COMPILE execution boundary. +# The analytical stage-5 estimate does not account for NATTEN's transient custom-op +# workspace, which can otherwise make a full-resolution tile allocate tens of GiB. +_CHUNKED_COMPILE_MAX_TILE_SIZE: tuple[int, int, int] = (80, 320, 320) + + +def stage5_mem_coef(mode: DiffVAEMode) -> float: + """Stage-5 working-set multiplicity for auto tiling, after host NA resolve. + ``CHUNKED_COMPILE``'s coef 7 assumes natten. When the host remaps chunked + modes to Triton/eager fallback, use the ``CHUNKED_EAGER`` coef (5) instead. + ``COMBINED_COMPILE`` requires natten and keeps coef 11. + """ + try: + base = _MEM_COEF_BY_MODE[mode] + except KeyError as exc: + raise ValueError(f"Unsupported DiffVAEMode for tiling budget: {mode!r}") from exc + return base + + +def budget_safety_bytes(mode: DiffVAEMode) -> int: + """Extra bytes withheld from the recommend budget (eager 1 GiB, compiled 2 GiB).""" + if mode is DiffVAEMode.CHUNKED_EAGER: + return _BUDGET_SAFETY_BYTES_EAGER + return _BUDGET_SAFETY_BYTES_COMPILED + + +def _max_tile_size(mode: DiffVAEMode) -> tuple[int, int, int] | None: + """Return the execution-mode-specific upper bound for an auto-selected tile.""" + if mode is DiffVAEMode.CHUNKED_COMPILE: + return _CHUNKED_COMPILE_MAX_TILE_SIZE + return None + + +def accumulator_element_size(feature_dtype: torch.dtype) -> int: + """Bytes per accumulator element; mirrors ``_decode_temporal_group_isolated``. + ``accum_dtype = float16 if feat_s4.dtype == bfloat16 else feat_s4.dtype``. + """ + if feature_dtype is torch.bfloat16: + return 2 # stored as fp16 + return int(torch.tensor([], dtype=feature_dtype).element_size()) + + +def stage4_feature_bytes( + *, + height: int, + width: int, + num_frames: int, + upsample_strides: Sequence[Tuple[int, int, int]], + stage4_channels: int, + element_size: int = _DEFAULT_ELEMENT_SIZE, + natten_trailing_pad_latent_frames: int = 0, +) -> int: + """Resident stages-1-3 output size (full volume tiled into stage 4). + Matches ``DiffusionVideoDecoder.forward_stages_1_to_3`` after optional NATTEN + trailing latent pad: channels-last ``(B, T, H, W, C)`` at stage-4 input resolution. + """ + if stage4_channels < 1: + raise ValueError(f"stage4_channels must be >= 1, got {stage4_channels}") + if element_size < 1: + raise ValueError(f"element_size must be >= 1, got {element_size}") + if len(upsample_strides) < 3: + raise ValueError(f"need at least 3 upsample strides, got {len(upsample_strides)}") + if natten_trailing_pad_latent_frames < 0: + raise ValueError(f"natten_trailing_pad_latent_frames must be >= 0, got {natten_trailing_pad_latent_frames}") + + # Local import: types ↔ tiling cycle avoidance at module import time. + from .types import VIDEO_SCALE_FACTORS, VideoLatentShape, VideoPixelShape # noqa: PLC0415 + + latent = VideoLatentShape.from_pixel_shape( + VideoPixelShape(batch=1, frames=int(num_frames), height=int(height), width=int(width), fps=24.0), + scale_factors=VIDEO_SCALE_FACTORS, + ) + s4_t, s4_h, s4_w = stage4_thw_from_latent( + upsample_strides[:3], + latent.frames + int(natten_trailing_pad_latent_frames), + latent.height, + latent.width, + drop_leading_frame=True, + ) + return int(s4_t) * int(s4_h) * int(s4_w) * int(stage4_channels) * int(element_size) + + +# --------------------------------------------------------------------------- +# Public entry points (pipeline recommend / decode schedule) +# --------------------------------------------------------------------------- + + +def recommended_decode_tiling_config( # noqa: PLR0913 + *, + tile_halos: Tuple[Tuple[int, int, int], Tuple[int, int, int]], + pixel_scale: SpatioTemporalScaleFactors, + min_tile_size_s4: Tuple[int, int, int], + patch_size: int, + height: int, + width: int, + num_frames: int, + mode: DiffVAEMode, + free_bytes: int, + stage5_channels: int, + stage4_channels: int, + upsample_strides: Sequence[Tuple[int, int, int]], + model_bytes: int = 0, + element_size: int = _DEFAULT_ELEMENT_SIZE, + natten_trailing_pad_latent_frames: int = 0, + out_channels: int = _ACCUMULATOR_CHANNELS, +) -> TileSizeConfig: + """Pick DiffVAE decode tiling from stage-4/5 halos and free VRAM. + Always enables both spatial and temporal tiling (temporal-only full-frame slabs + are unsafe on Hopper / some natten builds). + Selection (size-grid, accumulator-aware): + 1. Enumerate legal tile **sizes** on the LCM of DiffVAE ``pixel_scale`` and + :data:`~ltx_core.types.VIDEO_SCALE_FACTORS` (so configs also pass + :class:`~ltx_core.tiling.TileSizeConfig` construction); derive tile + counts from :func:`~ltx_core.tiling.split_by_size` (same as decode). + 2. Drop triples whose peak-bytes estimate exceeds ``usable`` bytes + (``free - max(model, 1 GiB) - safety - stage4_feature``; safety is + 1 GiB eager / 2 GiB compiled). Stage-4 input features stay resident + for the whole tiled decode. + 3. For ``CHUNKED_COMPILE``, cap each tile to 80 frames by 320x320 pixels. + This guards the NATTEN custom-op workspace not represented by the + analytical estimate. + 4. Among feasible triples, pick minimal :func:`volumetric_overlap_waste`. + Peak-bytes estimate:: + stage4_feature_bytes(...) # hard, full volume + + H * W * (2 * tile_t) * out_channels * element_size + + stage5_tokens * stage5_channels * element_size * coef + Accumulator is full output HxW (not spatially tiled) with temporal extent + ``2 * tile_t``: current group buffer plus the still-live previous exclusive + emit / overlap stub during handoff (not merely ``tile_t + overlap_t``). + RGBx``element_size`` by default. ``element_size`` is the activation width: + production bf16 features use fp16 accumulators (2), matching + :func:`accumulator_element_size`. Stage-5 uses the same element size x + ``stage5_channels`` x ``coef`` (11 / 7 / 5 / 2.5 by mode). + """ + if height < 1 or width < 1 or num_frames < 1: + raise ValueError(f"height/width/num_frames must be >= 1, got {height}x{width}x{num_frames}") + if patch_size < 1: + raise ValueError(f"patch_size must be >= 1, got {patch_size}") + if stage5_channels < 1: + raise ValueError(f"stage5_channels must be >= 1, got {stage5_channels}") + if out_channels < 1: + raise ValueError(f"out_channels must be >= 1, got {out_channels}") + if element_size < 1: + raise ValueError(f"element_size must be >= 1, got {element_size}") + + overlap_t, overlap_hw = recommended_pixel_overlaps(tile_halos, pixel_scale) + + ft, fh, fw = pixel_scale.time, pixel_scale.height, pixel_scale.width + # Construction validates fixed 8/32/32; to_splitters uses pixel_scale - step both. + step_t = math.lcm(ft, VIDEO_SCALE_FACTORS.time) + step_h = math.lcm(fh, VIDEO_SCALE_FACTORS.height) + step_w = math.lcm(fw, VIDEO_SCALE_FACTORS.width) + min_t_px = _round_up( + # ``2 * overlap`` so left+right ramps fit (else masks are not complementary and + # decode allocates a full weights buffer ≈ another accumulator). + max(2 * ft, 2 * overlap_t, _round_up(min_tile_size_s4[0] * ft, ft), 16), + step_t, + ) + min_h_px = _round_up( + max(2 * fh, 2 * overlap_hw, _round_up(min_tile_size_s4[1] * fh, fh), 64), + step_h, + ) + min_w_px = _round_up( + max(2 * fw, 2 * overlap_hw, _round_up(min_tile_size_s4[2] * fw, fw), 64), + step_w, + ) + + model_cost = max(int(model_bytes), _MIN_MODEL_BYTES_FLOOR) + coef = stage5_mem_coef(mode) + s4_feat_bytes = stage4_feature_bytes( + height=height, + width=width, + num_frames=num_frames, + upsample_strides=upsample_strides, + stage4_channels=stage4_channels, + element_size=element_size, + natten_trailing_pad_latent_frames=natten_trailing_pad_latent_frames, + ) + usable = max(0, int(free_bytes) - model_cost - budget_safety_bytes(mode) - s4_feat_bytes) + s5_bytes_per_token = max(1.0, float(stage5_channels) * float(element_size) * coef) + acc_bytes_per_pixel = int(out_channels) * int(element_size) + + t_cands = _axis_candidates(num_frames, overlap_t, min_t_px, step_t) + h_cands = _axis_candidates(height, overlap_hw, min_h_px, step_h) + w_cands = _axis_candidates(width, overlap_hw, min_w_px, step_w) + max_tile_size = _max_tile_size(mode) + if max_tile_size is not None: + t_cands = [candidate for candidate in t_cands if candidate[0] <= max_tile_size[0]] + h_cands = [candidate for candidate in h_cands if candidate[0] <= max_tile_size[1]] + w_cands = [candidate for candidate in w_cands if candidate[0] <= max_tile_size[2]] + + scored: list[tuple[float, int, int, int, int, int]] = [] + # (waste, -volume, n_t*n_h*n_w, tile_t, tile_h, tile_w) - minimize waste, then launches. + for tile_t, n_t in t_cands: + # Current group buffer + still-live emit/stub during temporal handoff. + acc_frames = 2 * int(tile_t) + acc_bytes = acc_frames * int(height) * int(width) * acc_bytes_per_pixel + if acc_bytes >= usable: + continue + s5_budget_bytes = usable - acc_bytes + max_s5_tokens = int(s5_budget_bytes // s5_bytes_per_token) + for tile_h, n_h in h_cands: + for tile_w, n_w in w_cands: + if stage5_tokens_for_pixel_tile(tile_t, tile_h, tile_w, patch_size=patch_size) > max_s5_tokens: + continue + waste = volumetric_overlap_waste( + num_frames=num_frames, + height=height, + width=width, + tile_frames=tile_t, + tile_height=tile_h, + tile_width=tile_w, + n_t=n_t, + n_h=n_h, + n_w=n_w, + ) + scored.append((waste, -tile_t * tile_h * tile_w, n_t * n_h * n_w, tile_t, tile_h, tile_w)) + + if not scored: + raise ValueError( + "Cannot fit a DiffVAE decode tile under the memory budget: " + f"min tile ~{min_t_px}f x {min_h_px}x{min_w_px}px " + f"(overlaps T={overlap_t}, HW={overlap_hw}), " + f"mode={mode.value}, coef={coef}, stage5_channels={stage5_channels}, " + f"stage4_feature_bytes={s4_feat_bytes}, usable_bytes={usable}. " + "Reduce resolution or free GPU memory." + ) + + scored.sort() + _waste, _vol, _ntiles, tile_t, tile_h, tile_w = scored[0] + return TileSizeConfig( + frames=DimensionSizeConfig(tile_size=tile_t, overlap=overlap_t), + height=DimensionSizeConfig(tile_size=tile_h, overlap=overlap_hw), + width=DimensionSizeConfig(tile_size=tile_w, overlap=overlap_hw), + ) + + +def prepare_tile_schedule( + stage4_shape_bcthw: torch.Size, + tiling_config: TilingConfig | None, + *, + upsample3_stride: Tuple[int, int, int], + patch_size: int, + min_tile_size: Tuple[int, int, int], + tile_halos: Tuple[Tuple[int, int, int], Tuple[int, int, int]], +) -> List[Tile]: + """Build pixel-blend tiles whose ``in_coords`` land on the stage-4 input grid. + DiffVAE temporal tiling deliberately skips ConvVAE causal split/mask tricks + (``split_temporal_causal``, ``left_starts_from_0``): pixel overlap already + covers blend+halo, and interval propagation follows + :class:`~ltx_core.model.video_vae.transformer.layers.LinearPixelShuffleUpsample` + (``drop_leading_frame`` only on the origin tile) with *symmetric* trapezoid + ramps so masks stay complementary without a weight buffer. + """ + pixel_scale = stage4_to_pixel_scale_factors(upsample3_stride, patch_size) + if tiling_config is None: + return [ + Tile( + in_coords=(slice(None), slice(None), slice(None), slice(None), slice(None)), + out_coords=(slice(None), slice(None), slice(None), slice(None), slice(None)), + masks_1d=( + untiled_mask_1d(), + untiled_mask_1d(), + untiled_mask_1d(), + untiled_mask_1d(), + untiled_mask_1d(), + ), + ) + ] + + overlap_t, overlap_hw = recommended_pixel_overlaps(tile_halos, pixel_scale) + _validate_overlap(tiling_config, min_overlap_frames=overlap_t, min_overlap_pixels=overlap_hw) + # Plain split (not split_temporal_causal): no start-1 / left_ramp+1 copycat of ConvVAE. + t_split, h_split, w_split = tiling_config.to_splitters( + pixel_scale, min_tile_size=min_tile_size, causal_temporal=False + ) + st, sh, sw = upsample3_stride + + def axis_specs( + split_op: SplitOperation, + dim_len: int, + stride_component: int, + *, + propagate_causal: bool, + apply_patch: bool, + ) -> list[tuple[slice, slice, torch.Tensor]]: + if split_op is DEFAULT_SPLIT_OPERATION: + return [(slice(None), slice(None), untiled_mask_1d())] + intervals = split_op(dim_len).intervals + specs = [] + for iv in intervals: + stage5 = _propagate_interval_through_upsample_hops(iv, [stride_component], propagate_causal) + if apply_patch: + pixel = _propagate_interval_through_upsample_hops(stage5, [patch_size], causal=False) + else: + pixel = stage5 + # Symmetric ramps (left_starts_from_0=False) for partition-of-unity with + # pixel-shuffle out_coords; ConvVAE sacrificial first-sample is not used. + mask_pixel = compute_trapezoidal_mask_1d( + pixel.end - pixel.start, pixel.left_ramp, pixel.right_ramp, left_starts_from_0=False + ) + specs.append((slice(iv.start, iv.end), slice(pixel.start, pixel.end), mask_pixel)) + return specs + + # Temporal: pixel-shuffle propagate (drop-leading geometry); spatial: exact x stride. + t_specs = axis_specs(t_split, stage4_shape_bcthw[2], st, propagate_causal=True, apply_patch=False) + h_specs = axis_specs(h_split, stage4_shape_bcthw[3], sh, propagate_causal=False, apply_patch=True) + w_specs = axis_specs(w_split, stage4_shape_bcthw[4], sw, propagate_causal=False, apply_patch=True) + + tiles: List[Tile] = [] + for t_spec, h_spec, w_spec in itertools.product(t_specs, h_specs, w_specs): + t_s4, t_px, t_mask = t_spec + h_s4, h_px, h_mask = h_spec + w_s4, w_px, w_mask = w_spec + tiles.append( + Tile( + in_coords=(slice(None), t_s4, h_s4, w_s4, slice(None)), + out_coords=(slice(None), slice(None), t_px, h_px, w_px), + masks_1d=(untiled_mask_1d(), untiled_mask_1d(), t_mask, h_mask, w_mask), + ) + ) + return tiles + + +def slice_stage4_tile( + feat_s4: torch.Tensor, + tile: Tile, + *, + content_frames: int, +) -> tuple[torch.Tensor, bool, bool, tuple[int, int, int]]: + """Slice a stage-4 feature tile, extending trailing tiles to include ghost frames.""" + is_origin = tile.in_coords[1].start in (0, None) + _, stop, _ = tile.in_coords[1].indices(content_frames) + pad_trailing = stop == content_frames + _b, t_coord, h_coord, w_coord, _c = tile.in_coords + t0, t1, _ = t_coord.indices(content_frames) + h0, h1, _ = h_coord.indices(feat_s4.shape[2]) + w0, w1, _ = w_coord.indices(feat_s4.shape[3]) + content_thw = (t1 - t0, h1 - h0, w1 - w0) + if pad_trailing: + t1 = feat_s4.shape[1] + feat_tile = feat_s4[:, t0:t1, h_coord, w_coord, :] + return feat_tile, is_origin, pad_trailing, content_thw + + +# --------------------------------------------------------------------------- +# Common helpers (geometry, pad/crop, stage floors) +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class AxisPad: + """How many elements were added (pad) or removed (crop) on each side of one axis.""" + + before: int + after: int + + +def resize_axis( + x: torch.Tensor, + dim: int, + size: int, + *, + mode: ResizeAxisMode, +) -> tuple[torch.Tensor, AxisPad]: + """Pad or crop axis ``dim`` so its length becomes ``size``. + Pad (``len < size``): + ``repeat_last`` - append copies of the last slice. + ``symmetric`` - edge-replicate first/last; leftover goes to the end + (``before = need // 2``, ``after = need - before``). + Crop (``len > size``): + ``repeat_last`` - drop from the end. + ``symmetric`` - drop from both ends with the same split rule as pad. + """ + if size < 1: + raise ValueError(f"resize_axis target size must be >= 1, got {size}") + if dim < 0: + dim += x.ndim + if not 0 <= dim < x.ndim: + raise ValueError(f"dim {dim} out of range for rank-{x.ndim} tensor") + + length = x.shape[dim] + if length == size: + return x, AxisPad(0, 0) + + if length < size: + need = size - length + if mode == "repeat_last": + last = x.narrow(dim, length - 1, 1) + expand_shape = list(x.shape) + expand_shape[dim] = need + pad = last.expand(expand_shape) + return torch.cat([x, pad], dim=dim), AxisPad(0, need) + + before = need // 2 + after = need - before + first = x.narrow(dim, 0, 1) + last = x.narrow(dim, length - 1, 1) + parts: list[torch.Tensor] = [] + if before: + expand_shape = list(x.shape) + expand_shape[dim] = before + parts.append(first.expand(expand_shape)) + parts.append(x) + if after: + expand_shape = list(x.shape) + expand_shape[dim] = after + parts.append(last.expand(expand_shape)) + return torch.cat(parts, dim=dim), AxisPad(before, after) + + need = length - size + if mode == "repeat_last": + return x.narrow(dim, 0, size).contiguous(), AxisPad(0, need) + + before = need // 2 + after = need - before + return x.narrow(dim, before, size).contiguous(), AxisPad(before, after) + + +def ensure_min_latent_shape( + latent: torch.Tensor, + min_tile_sizes: Tuple[int, int, int], +) -> tuple[torch.Tensor, tuple[AxisPad, AxisPad, AxisPad]]: + """Pad latent ``(B, C, T, H, W)`` up to ``min_tile_sizes`` if needed.""" + min_t, min_h, min_w = min_tile_sizes + t_pad = AxisPad(0, 0) + h_pad = AxisPad(0, 0) + w_pad = AxisPad(0, 0) + x = latent + if x.shape[2] < min_t: + x, t_pad = resize_axis(x, 2, min_t, mode="repeat_last") + if x.shape[3] < min_h: + x, h_pad = resize_axis(x, 3, min_h, mode="symmetric") + if x.shape[4] < min_w: + x, w_pad = resize_axis(x, 4, min_w, mode="symmetric") + return x, (t_pad, h_pad, w_pad) + + +def scale_axis_pad(pad: AxisPad, scale: int) -> AxisPad: + """Scale a latent-grid ``AxisPad`` into pixel (or other) units.""" + return AxisPad(pad.before * scale, pad.after * scale) + + +def crop_pixels_to_content( + pixels: torch.Tensor, + frames: int, + height: int, + width: int, + *, + h_pad: AxisPad | None = None, + w_pad: AxisPad | None = None, + spatial_scale: Tuple[int, int] = (1, 1), +) -> torch.Tensor: + """Crop padded decode output ``(B, C, F, H, W)`` back to the content shape. + Temporal pad is always trailing (``repeat_last``), so T is cropped from the + end. Spatial size-floor pads must pass the recorded ``h_pad`` / ``w_pad`` + (latent units) plus ``spatial_scale`` ``(H, W)`` so odd leftovers are not + re-split by a center-crop after upscaling. + """ + x, _ = resize_axis(pixels, 2, frames, mode="repeat_last") + scale_h, scale_w = spatial_scale + if h_pad is not None: + before = scale_axis_pad(h_pad, scale_h).before + if before + height > x.shape[3]: + raise ValueError(f"H crop out of range: before={before}, height={height}, got {x.shape[3]}") + x = x.narrow(3, before, height).contiguous() + else: + x, _ = resize_axis(x, 3, height, mode="symmetric") + if w_pad is not None: + before = scale_axis_pad(w_pad, scale_w).before + if before + width > x.shape[4]: + raise ValueError(f"W crop out of range: before={before}, width={width}, got {x.shape[4]}") + x = x.narrow(4, before, width).contiguous() + else: + x, _ = resize_axis(x, 4, width, mode="symmetric") + return x + + +def stage5_pixel_shape_from_stage4( + stage4_t: int, + stage4_h: int, + stage4_w: int, + *, + upsample_stride: Tuple[int, int, int], + patch_size: int, + stage5_kernel_t: int, + drop_leading_frame: bool, + pad_trailing: bool, +) -> tuple[int, int, int]: + """Pixel ``(F, H, W)`` for a stage-4-input extent (one remaining NA hop + patch).""" + st, sh, sw = upsample_stride + frames = stage4_t * st - 1 if drop_leading_frame and st == 2 else stage4_t * st + if pad_trailing: + frames = max(frames, stage5_kernel_t) + return frames, stage4_h * sh * patch_size, stage4_w * sw * patch_size + + +def pad_trailing_latent_for_natten_border(latent: torch.Tensor, n_frames: int) -> torch.Tensor: + """Replicate the last latent frame ``n_frames`` times for NATTEN last-frame border.""" + if n_frames <= 0: + return latent + padded, _ = resize_axis(latent, 2, latent.shape[2] + n_frames, mode="repeat_last") + return padded + + +def crop_trailing_context_natten_pad( + context: torch.Tensor, + *, + n_latent_frames: int, + time_scale: int, + stage5_kernel_t: int, +) -> torch.Tensor: + """Crop ghosting appendix before stage 5, leaving at least ``stage5_kernel_t``.""" + if n_latent_frames <= 0: + return context + ghost = n_latent_frames * time_scale + content_t = max(context.shape[1] - ghost, 1) + keep = min(context.shape[1], max(content_t, stage5_kernel_t)) + cropped, _ = resize_axis(context, 1, keep, mode="repeat_last") + return cropped + + +def _weight_floor(dtype: torch.dtype) -> float: + """Smallest divisor that safely guards ``buffer / weights`` in ``dtype``.""" + return max(1e-8, torch.finfo(dtype).tiny) + + +def stage4_thw_from_latent( + upsample_strides: Sequence[Tuple[int, int, int]], + latent_t: int, + latent_h: int, + latent_w: int, + *, + drop_leading_frame: bool = True, +) -> Tuple[int, int, int]: + """Stage-4 input ``(T, H, W)`` after the first three upsample hops.""" + t, h, w = latent_t, latent_h, latent_w + for st, sh, sw in upsample_strides[:3]: + t, h, w = t * st, h * sh, w * sw + if st == 2 and drop_leading_frame: + t -= 1 + return t, h, w + + +def stage4_to_pixel_scale_factors( + upsample_stride: Tuple[int, int, int], + patch_size: int, +) -> SpatioTemporalScaleFactors: + """Pixel/frame units per stage-4-input cell (last NA hop + unpatchify).""" + st, sh, sw = upsample_stride + return SpatioTemporalScaleFactors(time=st, height=sh * patch_size, width=sw * patch_size) + + +def compute_tile_min_size( + stage4_kernel: Tuple[int, int, int], + stage5_kernel: Tuple[int, int, int], + upsample3_stride: Tuple[int, int, int], +) -> Tuple[int, int, int]: + """Min stage-4-input ``(T, H, W)`` so stages 4 and 5 each see ``>= kernel``.""" + return tuple(max(stage4_kernel[a], -(-stage5_kernel[a] // upsample3_stride[a])) for a in range(3)) + + +def compute_tile_halos( + stage4_kernel: Tuple[int, int, int], + stage4_depth: int, + stage5_kernel: Tuple[int, int, int], + stage5_depth: int, + upsample3_stride: Tuple[int, int, int], +) -> Tuple[Tuple[int, int, int], Tuple[int, int, int]]: + """One-sided halos in stage-4-input units for stages 4 and 5.""" + halo4 = tuple(stage4_depth * (stage4_kernel[a] // 2) for a in range(3)) + halo5 = tuple(-(-(stage5_depth * (stage5_kernel[a] // 2)) // upsample3_stride[a]) for a in range(3)) + return halo4, halo5 # type: ignore[return-value] + + +def _cumulative_upsample_strides( + upsamples: Sequence[Tuple[Tuple[int, int, int], int]], +) -> List[Tuple[int, int, int]]: + """Per-axis product of hop strides for ``upsamples[:i]`` (``cumulative[0] = (1,1,1)``).""" + cumulative = [(1, 1, 1)] + t, h, w = 1, 1, 1 + for stride, _ in upsamples: + t, h, w = t * stride[0], h * stride[1], w * stride[2] + cumulative.append((t, h, w)) + return cumulative + + +def all_stages_min_tile_size( + stage_kernels: Sequence[Tuple[int, int, int]], + upsamples: Sequence[Tuple[Tuple[int, int, int], int]], + stage5_kernel: Tuple[int, int, int], +) -> Tuple[int, int, int]: + """Per-axis latent-grid floor so every stage's NA sees dims ``>= kernel_size``.""" + cumulative = _cumulative_upsample_strides(upsamples) + mins = [1, 1, 1] + for stage_i in range(len(upsamples)): + strides = cumulative[stage_i] + for axis in range(3): + mins[axis] = max(mins[axis], -(-stage_kernels[stage_i][axis] // strides[axis])) + strides5 = cumulative[len(upsamples)] + for axis in range(3): + mins[axis] = max(mins[axis], -(-stage5_kernel[axis] // strides5[axis])) + return (mins[0], mins[1], mins[2]) + + +def pixel_tile_shape(full_shape: tuple[int, ...], out_coords: tuple[slice, ...]) -> tuple[int, ...]: + dims: list[int] = [] + for size, coord in zip(full_shape, out_coords, strict=True): + start, stop, step = coord.indices(size) + dims.append(len(range(start, stop, step))) + return tuple(dims) + + +# --------------------------------------------------------------------------- +# Recommendation helpers +# --------------------------------------------------------------------------- + + +def _round_up(value: int, multiple: int) -> int: + return -(-value // multiple) * multiple + + +def recommended_pixel_overlaps( + tile_halos: Tuple[Tuple[int, int, int], Tuple[int, int, int]], + pixel_scale: SpatioTemporalScaleFactors, +) -> Tuple[int, int]: + """Stage-4/5-safe ``(temporal_overlap_frames, spatial_overlap_pixels)``. + Shared by :func:`recommended_decode_tiling_config` (to *set* overlaps) and + :func:`~ltx_core.tiling._validate_overlap` (to reject undersized configs). + """ + + def dominant(axis: int) -> int: + return max(tile_halos[i][axis] for i in range(len(tile_halos))) + + overlap_t = _round_up(dominant(0) * pixel_scale.time, 8) + halo_hw = max(dominant(1), dominant(2)) + overlap_hw = _round_up(halo_hw * pixel_scale.height, 32) + return overlap_t, overlap_hw + + +def stage5_tokens_for_pixel_tile( + tile_frames: int, + tile_height: int, + tile_width: int, + *, + patch_size: int, +) -> int: + """Pre-unpatchify stage-5 token count for a pixel-space tile (NATTEN volume).""" + h5 = max(1, tile_height // patch_size) + w5 = max(1, tile_width // patch_size) + return tile_frames * h5 * w5 + + +def _axis_candidates(length: int, overlap: int, min_size: int, multiple: int) -> list[tuple[int, int]]: + """``(tile_size, num_tiles)`` for every legal size on ``multiple``'s grid.""" + out: list[tuple[int, int]] = [] + max_size = max(_round_up(length, multiple), min_size) + for size in range(min_size, max_size + multiple, multiple): + if size <= overlap: + continue + n = len(split_by_size(size, overlap)(length).intervals) + out.append((size, n)) + return out + + +def volumetric_overlap_waste( + *, + num_frames: int, + height: int, + width: int, + tile_frames: int, + tile_height: int, + tile_width: int, + n_t: int, + n_h: int, + n_w: int, +) -> float: + """``processed_volume / unique_volume`` (>= 1). Lower means less overlap recompute.""" + processed = n_t * n_h * n_w * tile_frames * tile_height * tile_width + unique = max(1, num_frames * height * width) + return processed / unique + + +# --------------------------------------------------------------------------- +# Schedule helpers +# --------------------------------------------------------------------------- + + +def _propagate_interval_through_upsample_hops( + interval: DimensionInterval, + strides: Sequence[int], + causal: bool, +) -> DimensionInterval: + """Forward-propagate one interval through a sequence of upsample hops on one axis. + Mirrors :class:`~ltx_core.model.video_vae.transformer.layers.LinearPixelShuffleUpsample`: + multiply by ``stride``, and for the causal temporal axis when ``stride == 2`` apply + the duplicate-frame drop (``end -= 1``; non-origin also ``start -= 1``). + This is *not* :func:`~ltx_core.model.video_vae.video_vae.map_temporal_slice` (ConvVAE). + DiffVAE non-origin tiles run with ``drop_leading_frame=False`` and must keep length + ``tile_t * stride``; the ConvVAE ``1+(L-1)*stride`` mapping is one frame short and + shifts non-origin ``out_coords``, which breaks tiled↔untiled temporal blend even + when masks are complementary. + """ + x = interval + for stride in strides: + if stride < 1: + raise ValueError(f"upsample stride must be >= 1, got {stride}") + start = x.start * stride + end = x.end * stride + left_ramp = x.left_ramp * stride + right_ramp = x.right_ramp * stride + if causal and stride == 2: + end -= 1 + if x.start != 0: + start -= 1 + x = DimensionInterval(start=start, end=end, left_ramp=left_ramp, right_ramp=right_ramp) + return x diff --git a/telefuser/models/ltx25/diff_vae/diffusion_video_decoder.py b/telefuser/models/ltx25/diff_vae/diffusion_video_decoder.py new file mode 100644 index 0000000..76cffee --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/diffusion_video_decoder.py @@ -0,0 +1,912 @@ +"""Diffusion (NATTEN) video VAE decoder.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Iterator, List, Literal, Tuple + +import torch +from safetensors import safe_open +from torch import nn + +from telefuser.models.ltx25.diff_vae.ops import PerChannelStatistics, patchify, unpatchify +from telefuser.models.ltx25.diff_vae.transformer import ( + AdaLNZero, + ChannelLinear, + CombinedDiffusionNABlock, + LinearPixelShuffleUpsample, + NABlock, +) +from telefuser.models.ltx25.diff_vae.transformer.attention import ( + NattenAttention, + NeighborhoodAttention3D, + configure_w_chunks, + natten_available, +) +from telefuser.models.ltx25.diff_vae.transformer.blocks import DiffusionNABlock +from telefuser.models.ltx25.diff_vae.transformer.chunked.block import ChunkedDiffusionNABlock +from telefuser.models.ltx25.diff_vae.transformer.compiling import compile_diffusion_decoder +from telefuser.models.ltx25.diff_vae.transformer.config import DiffVAEMode +from telefuser.models.ltx25.diff_vae.transformer.fallback_na import fallback_na_attention +from telefuser.models.ltx25.diff_vae.transformer.rope_math import rope_inv_freqs +from telefuser.models.ltx25.transformer import PixArtAlphaCombinedTimestepSizeEmbeddings + +from ..checkpoint import inspect_checkpoint +from . import diffusion_tiling +from .tiling import ( + Tile, + TileSizeConfig, + TilingConfig, + group_tiles_by_temporal_slice, + masks_are_complementary, + scale_by_masks_1d, +) +from .types import SpatioTemporalScaleFactors, VideoLatentShape + + +def cuda_activation_budget_bytes(device: torch.device) -> int: + """Return CUDA bytes available for optional decoder tiling.""" + if device.type != "cuda": + return 0 + index = device.index if device.index is not None else torch.cuda.current_device() + free_bytes, _ = torch.cuda.mem_get_info(index) + return int(free_bytes) + + +def to_velocity(sample: torch.Tensor, sigma: float | torch.Tensor, denoised_sample: torch.Tensor) -> torch.Tensor: + """Convert an x0 prediction into the velocity parameterization.""" + sigma_value = sigma.to(torch.float32).item() if isinstance(sigma, torch.Tensor) else sigma + if sigma_value == 0: + raise ValueError("Sigma cannot be zero") + return ((sample.float() - denoised_sample.float()) / sigma_value).to(sample.dtype) + + +# Production CausalVideoAutoencoderL decoder layout: channel width per stage. +# Mirrors the (non-diffusion) NA decoder's stage spec. +_L_STAGE_CHANNELS: Tuple[int, ...] = (1024, 512, 256, 256, 128) +_L_STAGE_DEPTHS: Tuple[int, ...] = (4, 6, 4, 2, 2) +# (stride, out_channels_reduction_factor) per upsample, in stage order. +_L_UPSAMPLES: Tuple[Tuple[Tuple[int, int, int], int], ...] = ( + ((1, 2, 2), 2), # compress_space x2 + ((2, 1, 1), 2), # compress_time x2 + ((2, 2, 2), 1), # compress_all x1 (channel-preserving) + ((2, 2, 2), 2), # compress_all x2 +) +# Per-stage 3D neighborhood (K_t, K_h, K_w). +_L_STAGE_KERNELS: Tuple[Tuple[int, int, int], ...] = ( + (3, 7, 7), + (3, 7, 7), + (3, 5, 5), + (3, 5, 5), + (3, 3, 3), +) + +# Stage-5 (diffusion stage) defaults: wider kernel + more blocks than the +# deterministic stages, since it carries the entire per-step diffusion compute. +_DIFF_STAGE5_KERNEL_DEFAULT: Tuple[int, int, int] = (3, 7, 7) +_DIFF_STAGE5_DEPTH_DEFAULT: int = 8 +_DIFF_STAGE_DEPTHS_DEFAULT: Tuple[int, ...] = (*_L_STAGE_DEPTHS[:-1], _DIFF_STAGE5_DEPTH_DEFAULT) + + +class DiffusionVideoDecoder(nn.Module): + """Diffusion-based video VAE decoder (Neighborhood-Attention backbone). + Minimal port of the reference ``NADiffusionDecoder``. + Stages 1-4 deterministically upsample the latent into a context volume + (same NA-upsample path as the non-diffusion NA decoder). Stage 5 runs + ``DiffusionNABlock``s that denoise the patchified noised pixels ``x_t``, + guided by that context via AdaLN-Zero scale/shift (ungated residuals; + legacy static gates are folded into Linear weights at load time). + Last-frame NATTEN window-shift is mitigated by temporarily replicating the + last latent frame ``(stage1_K_t // 2) * 2`` times through stages 1-4, then + cropping that appendix from context before stage 5 - but only down to + ``max(original_context_T, stage5_kernel[0])`` so undersized clips (e.g. a + single latent frame) still satisfy NATTEN's kernel floor. Latents / tiles + below ``stage_min_tile_sizes`` are edge-padded first via ``diffusion_tiling``; + leftover pad is cropped from the final pixels. + """ + + def __init__( # noqa: PLR0913 + self, + in_channels: int = 128, + out_channels: int = 3, + patch_size: int = 4, + head_dim: int = 64, + rope_dim_split: Tuple[int, int, int] | None = None, + stage_channels: Tuple[int, ...] = _L_STAGE_CHANNELS, + stage_depths: Tuple[int, ...] = _DIFF_STAGE_DEPTHS_DEFAULT, + stage_kernels: Tuple[Tuple[int, int, int], ...] = _L_STAGE_KERNELS, + upsamples: Tuple[Tuple[Tuple[int, int, int], int], ...] = _L_UPSAMPLES, + stage5_kernel: Tuple[int, int, int] = _DIFF_STAGE5_KERNEL_DEFAULT, + stage5_channels: int | None = None, + t_emb_dim: int = 384, + default_num_inference_steps: int = 2, + timestep_scale_multiplier: float = 1.0, + model_output_type: Literal["v", "x0"] = "v", + ) -> None: + super().__init__() + assert len(stage_channels) == len(stage_depths) == len(stage_kernels) + assert len(upsamples) == len(stage_channels) - 1 + for c in stage_channels: + assert c % head_dim == 0, f"stage_channels {stage_channels} must each be a multiple of head_dim={head_dim}" + + self.patch_size = patch_size + self.register_buffer( + "default_inference_timesteps", + torch.linspace(1.0, 1.0 / default_num_inference_steps, default_num_inference_steps, device="cpu"), + persistent=False, + ) + self.out_channels = out_channels + self.stage_channels = stage_channels + self.stage_depths = stage_depths + self.base_channels = stage_channels[-1] + self.causal = False + self.timestep_conditioning = True + self.video_downscale_factors = SpatioTemporalScaleFactors.default() + self.stage5_kernel: Tuple[int, int, int] = tuple(stage5_kernel) # type: ignore[assignment] + # NATTEN last-frame border workaround: replicate last latent frame + # ``(K_t // 2) * 2`` times through stages 1-4, then crop the appendix + # off context before stage 5 down to at least ``stage5_kernel[0]``. + self._natten_trailing_pad_latent_frames = (stage_kernels[0][0] // 2) * 2 + + # Encoder output is per-channel normalized; undo before conv_in (same as ConvVideoDecoder). + self.per_channel_statistics = PerChannelStatistics(latent_channels=in_channels) + + self.conv_in = ChannelLinear(in_channels, stage_channels[0], bias=True) + + self.det_stages = nn.ModuleList() + self.upsamples = nn.ModuleList() + n_det_stages = len(stage_channels) - 1 + for stage_i in range(n_det_stages): + c = stage_channels[stage_i] + depth = stage_depths[stage_i] + kernel = stage_kernels[stage_i] + self.det_stages.append( + nn.ModuleList( + [ + NABlock(dim=c, kernel_size=kernel, head_dim=head_dim, rope_dim_split=rope_dim_split) + for _ in range(depth) + ] + ) + ) + stride, reduction = upsamples[stage_i] + self.upsamples.append( + LinearPixelShuffleUpsample(in_channels=c, stride=stride, out_channels_reduction_factor=reduction) + ) + + self.t_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim=t_emb_dim, size_emb_dim=0) + + c_ctx = stage_channels[-1] + self.context_channels = c_ctx + c5 = stage5_channels if stage5_channels is not None else c_ctx + self.stage5_channels = c5 + d5 = stage_depths[-1] + assert c5 % head_dim == 0, f"stage5_channels {c5} must be a multiple of head_dim={head_dim}" + noised_pixel_channels = out_channels * (patch_size**2) + + # Latent-grid floor so stages 1-3 (full volume) never undershoot NA. + self.stage_min_tile_sizes: Tuple[int, int, int] = diffusion_tiling.all_stages_min_tile_size( + stage_kernels, upsamples, stage5_kernel + ) + # Stage-4-input tile floor / overlap halos (only stages 4-5 are tiled). + up3_stride = upsamples[3][0] + self.tile_min_sizes: Tuple[int, int, int] = diffusion_tiling.compute_tile_min_size( + stage_kernels[3], stage5_kernel, up3_stride + ) + self.tile_halos: Tuple[Tuple[int, int, int], Tuple[int, int, int]] = diffusion_tiling.compute_tile_halos( + stage_kernels[3], + stage_depths[3], + stage5_kernel, + stage_depths[-1], + up3_stride, + ) + self.conv_in_x_t = ChannelLinear(noised_pixel_channels, c5, bias=True) + + # Shared AdaLN-Zero (7-chunk for shape compat; gate slots unused in block). + self.shared_adaln = AdaLNZero(dim=c5, t_emb_dim=t_emb_dim) + + self.diff_blocks = nn.ModuleList( + [ + CombinedDiffusionNABlock( + dim=c5, + kernel_size=stage5_kernel, + context_channels=c_ctx, + head_dim=head_dim, + rope_dim_split=rope_dim_split, + ) + for _ in range(d5) + ] + ) + + self.norm_out = nn.RMSNorm(c5, eps=1e-6) + self.conv_out = ChannelLinear(c5, noised_pixel_channels, bias=True) + + self.timestep_scale_multiplier = timestep_scale_multiplier + self.model_output_type = model_output_type + # Set True by ``compile_diffusion_decoder`` so decode marks T/H/W dynamic. + self.mark_dynamic_shapes = False + # When True, skip stage-4 upsample and inject via deferred sequential upsample+proj. + # Default False = combined pathway (``CombinedDiffusionNABlock``). Chunked DiffVAE + # modes flip this via ``apply_diffvae_config``. + self.deferred_stage4_upsample = False + + def _run_det_stage(self, x: torch.Tensor, stage_i: int, drop_leading_frame: bool) -> torch.Tensor: + """One deterministic stage: NA blocks + upsample.""" + if self.mark_dynamic_shapes: + for dim in (1, 2, 3): + torch._dynamo.mark_dynamic(x, dim) + for block in self.det_stages[stage_i]: + x = block(x) + return self.upsamples[stage_i](x, drop_leading_frame=drop_leading_frame) + + def forward_stages_1_to_3( + self, + z_noisy: torch.Tensor, + drop_leading_frame: bool = True, + ) -> torch.Tensor: + """Stages 1-3 on a full (or already ghost-padded) latent → stage-4 input feature. + Output is channels-last ``(B, T, H, W, C)`` at stage-4 input resolution. + Callers that want NATTEN trailing ghosting should pad the latent first via + ``diffusion_tiling.pad_trailing_latent_for_natten_border``. + """ + z_noisy = self.per_channel_statistics.un_normalize(z_noisy) + x = z_noisy.permute(0, 2, 3, 4, 1) + x = self.conv_in(x) + for stage_i in range(3): + x = self._run_det_stage(x, stage_i, drop_leading_frame) + return x + + def forward_stage_4( + self, + x: torch.Tensor, + drop_leading_frame: bool = True, + pad_trailing: bool = True, + ) -> torch.Tensor: + """Stage 4 on a stage-4-input feature tile → stage-5 context (or pre-upsample feat). + ``x`` is channels-last. When ``pad_trailing``, soft-crop the ghosting + appendix before returning (ghost pad must already be present upstream). + When ``deferred_stage4_upsample`` is set, runs NA blocks only (no + ``upsamples[3]``) and crops ghost at pre-upsample temporal resolution. + """ + if self.deferred_stage4_upsample: + if self.mark_dynamic_shapes: + for dim in (1, 2, 3): + torch._dynamo.mark_dynamic(x, dim) + for block in self.det_stages[3]: + x = block(x) + if pad_trailing: + up_t = int(self.upsamples[3].stride[0]) + x = diffusion_tiling.crop_trailing_context_natten_pad( + x, + n_latent_frames=self._natten_trailing_pad_latent_frames, + time_scale=self.video_downscale_factors.time // up_t, + stage5_kernel_t=max(1, -(-self.stage5_kernel[0] // up_t)), + ) + return x + + x = self._run_det_stage(x, 3, drop_leading_frame) + if pad_trailing: + x = diffusion_tiling.crop_trailing_context_natten_pad( + x, + n_latent_frames=self._natten_trailing_pad_latent_frames, + time_scale=self.video_downscale_factors.time, + stage5_kernel_t=self.stage5_kernel[0], + ) + return x + + def _context_and_x_for_diff_step(self, context: torch.Tensor, x_t: torch.Tensor) -> torch.Tensor: + """Build block-ready ``[context | conv_in_x_t(patched x)]`` for ``forward_diff_step``.""" + noised_pixels_patched = patchify(x_t, patch_size_hw=self.patch_size, patch_size_t=1) + x = self.conv_in_x_t(noised_pixels_patched.permute(0, 2, 3, 4, 1)) + return torch.cat([context, x], dim=-1) + + def _x_for_diff_step(self, x_t: torch.Tensor) -> torch.Tensor: + """Conv-processed noised pixels only (deferred-context path).""" + noised_pixels_patched = patchify(x_t, patch_size_hw=self.patch_size, patch_size_t=1) + return self.conv_in_x_t(noised_pixels_patched.permute(0, 2, 3, 4, 1)) + + def forward_diff_step( + self, + context_and_x: torch.Tensor, + t: torch.Tensor, + ) -> torch.Tensor: + """One stage-5 diffusion step. Returns the model prediction in pixel space. + ``context_and_x`` is ``[latent_context | conv_in_x_t(x)]`` (channels-last), built + at the call site via ``_context_and_x_for_diff_step``. That single buffer is + reused across ``diff_blocks``: each block writes its output x-half back + with ``copy_`` (no per-block ``cat``). One-tensor layout keeps Dynamo + T/H/W symbols identical under ``mark_dynamic``. + """ + x_half = context_and_x[..., self.context_channels :] + t_emb = self.t_embedder(self.timestep_scale_multiplier * t, hidden_dtype=x_half.dtype) + modulation = self.shared_adaln(t_emb) + + if self.mark_dynamic_shapes: + for dim in (1, 2, 3): + torch._dynamo.mark_dynamic(context_and_x, dim) + + for block in self.diff_blocks: + x_half.copy_(block.forward_combined(context_and_x, modulation)) + x = x_half + + x = self.norm_out(x) + x = self.conv_out(x) + x = x.permute(0, 4, 1, 2, 3).contiguous() + return unpatchify(x, patch_size_hw=self.patch_size, patch_size_t=1) + + def forward_diff_step_deferred( + self, + x: torch.Tensor, + stage4_feat: torch.Tensor, + t: torch.Tensor, + *, + drop_leading_frame: bool = True, + ) -> torch.Tensor: + """Stage-5 step with deferred context: only ``x`` + low-res ``stage4_feat``. + Marks T/H/W dynamic on both tensors. CHUNKED blocks upsample then + ``context_proj`` on the host before attn+mlp; BLACKWELL_DSL + (``DSLDiffusionBlockChain``) folds that hop into the fused kernel and never + materialises full-resolution context. ``drop_leading_frame`` must match the + flag used for this tile's stage-4 path (origin tile vs non-origin). + """ + t_emb = self.t_embedder(self.timestep_scale_multiplier * t, hidden_dtype=x.dtype) + modulation = self.shared_adaln(t_emb) + + if self.mark_dynamic_shapes: + for dim in (1, 2, 3): + torch._dynamo.mark_dynamic(x, dim) + torch._dynamo.mark_dynamic(stage4_feat, dim) + + for block in self.diff_blocks: + x = block.forward_x_ctx(x, stage4_feat, modulation, drop_leading_frame=drop_leading_frame) + + x = self.norm_out(x) + x = self.conv_out(x) + x = x.permute(0, 4, 1, 2, 3).contiguous() + return unpatchify(x, patch_size_hw=self.patch_size, patch_size_t=1) + + def _euler_step( + self, x_t: torch.Tensor, model_out: torch.Tensor, t_now: torch.Tensor, t_next: torch.Tensor + ) -> torch.Tensor: + """One reverse-diffusion Euler update: advance ``x_t`` from ``t_now`` to + ``t_next`` given the model's prediction at ``t_now``. + """ + compute_dtype = x_t.dtype + dt = (t_now - t_next).view(-1, *([1] * (x_t.ndim - 1))).to(torch.float32) + x_t_fp32 = x_t.to(torch.float32) + v_pred = model_out if self.model_output_type == "v" else to_velocity(x_t_fp32, t_now, model_out) + return (x_t_fp32 - dt * v_pred).to(compute_dtype) + + def _decode_one_tile( + self, + feat_tile: torch.Tensor, + x_t_tile_init: torch.Tensor, + *, + is_origin: bool, + timestep: torch.Tensor, + pad_trailing: bool, + ) -> torch.Tensor: + """Run stage 4 + diffusion on one stage-4 feature tile (isolation).""" + context_tile = self.forward_stage_4( + feat_tile, + drop_leading_frame=is_origin, + pad_trailing=pad_trailing, + ) + + x_t = x_t_tile_init + _, num_steps = timestep.shape + for i in range(num_steps - 1): + t_now = timestep[:, i] + t_next = timestep[:, i + 1] + if self.deferred_stage4_upsample: + x = self._x_for_diff_step(x_t) + model_out = self.forward_diff_step_deferred(x, context_tile, t_now, drop_leading_frame=is_origin).to( + torch.float32 + ) + else: + context_and_x = self._context_and_x_for_diff_step(context_tile, x_t) + model_out = self.forward_diff_step(context_and_x, t_now).to(torch.float32) + x_t = self._euler_step(x_t, model_out, t_now, t_next) + + t_now = timestep[:, -1] + if self.deferred_stage4_upsample: + x = self._x_for_diff_step(x_t) + model_out = self.forward_diff_step_deferred(x, context_tile, t_now, drop_leading_frame=is_origin) + else: + context_and_x = self._context_and_x_for_diff_step(context_tile, x_t) + model_out = self.forward_diff_step(context_and_x, t_now) + if self.model_output_type == "x0": + return model_out + return self._euler_step(x_t, model_out.to(torch.float32), t_now, torch.zeros_like(t_now)) + + def _decode_temporal_group_isolated( + self, + tiles: List[Tile], + feat_s4: torch.Tensor, + content_s4_frames: int, + x_t_init: torch.Tensor | None, + timestep: torch.Tensor, + full_video_shape: VideoLatentShape, + curr_temporal_slice: slice, + generator: torch.Generator | None, + *, + complementary: bool, + ) -> Tuple[torch.Tensor, torch.Tensor | None]: + """Decode every tile of one temporal group in isolation and blend.""" + group_temporal_len = curr_temporal_slice.stop - curr_temporal_slice.start + group_shape = full_video_shape._replace(frames=group_temporal_len) + full_torch_shape = full_video_shape.to_torch_shape() + accum_dtype = torch.float16 if feat_s4.dtype == torch.bfloat16 else feat_s4.dtype + buffer = torch.zeros(group_shape.to_torch_shape(), device=feat_s4.device, dtype=accum_dtype) + weights: torch.Tensor | None = None if complementary else torch.zeros_like(buffer) + local_temporal_slice = slice(0, group_temporal_len) + + compute_dtype = feat_s4.dtype + randn_device = generator.device if generator is not None else feat_s4.device + up3_stride = tuple(self.upsamples[3].stride) + + for tile in tiles: + feat_tile, is_origin, pad_trailing, content_thw = diffusion_tiling.slice_stage4_tile( + feat_s4, tile, content_frames=content_s4_frames + ) + content_pixel_shape = diffusion_tiling.pixel_tile_shape(full_torch_shape, tile.out_coords) + stage5_f, stage5_h, stage5_w = diffusion_tiling.stage5_pixel_shape_from_stage4( + content_thw[0], + content_thw[1], + content_thw[2], + upsample_stride=up3_stride, # type: ignore[arg-type] + patch_size=self.patch_size, + stage5_kernel_t=self.stage5_kernel[0], + drop_leading_frame=is_origin, + pad_trailing=pad_trailing, + ) + + if x_t_init is None: + x_t_tile_init = torch.randn( + (content_pixel_shape[0], content_pixel_shape[1], stage5_f, stage5_h, stage5_w), + dtype=compute_dtype, + generator=generator, + device=randn_device, + ).to(feat_s4.device) + else: + # Expand/crop to stage-5 canvas with the same edge policy as latent + # size-floor / ghost pad (not fresh noise - NA mixes padded values + # into kept pixels near the boundary). + x_t_tile_init = x_t_init[tile.out_coords] + x_t_tile_init, _ = diffusion_tiling.resize_axis(x_t_tile_init, 2, stage5_f, mode="repeat_last") + x_t_tile_init, _ = diffusion_tiling.resize_axis(x_t_tile_init, 3, stage5_h, mode="symmetric") + x_t_tile_init, _ = diffusion_tiling.resize_axis(x_t_tile_init, 4, stage5_w, mode="symmetric") + + pixel_tile = self._decode_one_tile( + feat_tile, + x_t_tile_init, + is_origin=is_origin, + timestep=timestep, + pad_trailing=pad_trailing, + ) + pixel_tile = diffusion_tiling.crop_pixels_to_content( + pixel_tile, + content_pixel_shape[2], + content_pixel_shape[3], + content_pixel_shape[4], + ).to(buffer.dtype) + + masks = tuple(m.to(device=buffer.device, dtype=torch.float32) for m in tile.masks_1d) + local_coords = ( + tile.out_coords[0], + tile.out_coords[1], + local_temporal_slice, + tile.out_coords[3], + tile.out_coords[4], + ) + buffer[local_coords] += scale_by_masks_1d(pixel_tile, masks) + if weights is not None: + strength = torch.ones(pixel_tile.shape, device=buffer.device, dtype=buffer.dtype) + weights[local_coords] += scale_by_masks_1d(strength, masks) + + return buffer, weights + + def _decode_pixels( # noqa: PLR0912, PLR0915 + self, + latent: torch.Tensor, + tiling_config: TilingConfig | None = None, + generator: torch.Generator | None = None, + *, + as_fhwc: bool = False, + ) -> Iterator[torch.Tensor]: + """Decode latent to pixels, yielding temporal chunks. + Default yields raw ``(B, C, F, H, W)`` in ``[-1, 1]``. With ``as_fhwc=True`` + (used by :meth:`decode_video`), each chunk is materialized once as + contiguous ``[F, H, W, C]`` still in ``[-1, 1]`` - layout copy only; + range mapping stays in ``to_rgb``. + Stages 1-3 run once on the full volume; stages 4-5 run per tile with + pixel blend (one tile / one group when untiled or no real split). + Across temporal groups only the trailing overlap is retained between + iterations; exclusive frames are yielded before the next group decodes. + Peak residency is ~two tile extents (current buffer + still-live emit / + overlap stub), not a single ``tile + overlap`` slab. + """ + content_shape = VideoLatentShape.from_torch_shape(latent.shape) + content_pixel = content_shape.upscale(self.video_downscale_factors)._replace(channels=self.out_channels) + + latent, (_t_pad, h_pad, w_pad) = diffusion_tiling.ensure_min_latent_shape(latent, self.stage_min_tile_sizes) + spatial_scale = (self.video_downscale_factors.height, self.video_downscale_factors.width) + work_shape = VideoLatentShape.from_torch_shape(latent.shape) + full_video_shape = work_shape.upscale(self.video_downscale_factors)._replace(channels=self.out_channels) + target_shape = full_video_shape.to_torch_shape() + + strides = [tuple(u.stride) for u in self.upsamples] + s4_t, s4_h, s4_w = diffusion_tiling.stage4_thw_from_latent( + strides, latent.shape[2], latent.shape[3], latent.shape[4], drop_leading_frame=True + ) + tiles = diffusion_tiling.prepare_tile_schedule( + torch.Size([latent.shape[0], latent.shape[1], s4_t, s4_h, s4_w]), + tiling_config, + upsample3_stride=tuple(self.upsamples[3].stride), # type: ignore[arg-type] + patch_size=self.patch_size, + min_tile_size=self.tile_min_sizes, + tile_halos=self.tile_halos, + ) + + latent_padded = diffusion_tiling.pad_trailing_latent_for_natten_border( + latent, self._natten_trailing_pad_latent_frames + ) + if self.mark_dynamic_shapes: + for dim in (2, 3, 4): + torch._dynamo.mark_dynamic(latent_padded, dim) + + feat_s4 = self.forward_stages_1_to_3(latent_padded, drop_leading_frame=True) + + batch = latent.shape[0] + timestep = self.default_inference_timesteps.to(latent.device).unsqueeze(0).expand(batch, -1) + single_step_x0 = timestep.shape[1] == 1 and self.model_output_type == "x0" + + x_t_init: torch.Tensor | None = None + if not single_step_x0: + compute_dtype = latent.dtype + randn_device = generator.device if generator is not None else latent.device + x_t_init = torch.randn( + tuple(target_shape), dtype=compute_dtype, generator=generator, device=randn_device + ).to(latent.device) + + complementary = masks_are_complementary(tiles, target_shape) + groups = group_tiles_by_temporal_slice(tiles) + group_slices = [slice(*group[0].out_coords[2].indices(target_shape[2])[:2]) for group in groups] + + # Keep only the trailing temporal overlap of the previous group (not the full + # chunk). Exclusive frames are yielded before the next group is decoded; the + # consumer may still hold that emit while the next buffer is live (~2x tile). + overlap_stub: torch.Tensor | None = None + overlap_stub_weights: torch.Tensor | None = None + + def _finalize(buf: torch.Tensor, wts: torch.Tensor | None) -> torch.Tensor: + if complementary: + return buf.to(latent.dtype) + assert wts is not None + wts = wts.clamp(min=diffusion_tiling._weight_floor(wts.dtype)) + return (buf / wts).to(latent.dtype) + + def _narrow_content_cfhw(t: torch.Tensor, frames_keep: int) -> torch.Tensor: + """Spatial/temporal content crop as views (no ``.contiguous()``).""" + x = t[:, :, :frames_keep] + th, tw = content_pixel.height, content_pixel.width + scale_h, scale_w = spatial_scale + if h_pad is not None: + before = diffusion_tiling.scale_axis_pad(h_pad, scale_h).before + x = x.narrow(3, before, th) + else: + need = x.shape[3] - th + if need > 0: + x = x.narrow(3, need // 2, th) + elif need < 0: + x, _ = diffusion_tiling.resize_axis(x, 3, th, mode="symmetric") + if w_pad is not None: + before = diffusion_tiling.scale_axis_pad(w_pad, scale_w).before + x = x.narrow(4, before, tw) + else: + need = x.shape[4] - tw + if need > 0: + x = x.narrow(4, need // 2, tw) + elif need < 0: + x, _ = diffusion_tiling.resize_axis(x, 4, tw, mode="symmetric") + return x + + def _crop_emit(buf: torch.Tensor, wts: torch.Tensor | None, global_start: int) -> torch.Tensor | None: + if global_start >= content_pixel.frames or buf.shape[2] < 1: + return None + frames_keep = min(buf.shape[2], content_pixel.frames - global_start) + if frames_keep < 1: + return None + if not as_fhwc: + chunk = _finalize(buf[:, :, :frames_keep], None if wts is None else wts[:, :, :frames_keep]) + return diffusion_tiling.crop_pixels_to_content( + chunk, + frames_keep, + content_pixel.height, + content_pixel.width, + h_pad=h_pad, + w_pad=w_pad, + spatial_scale=spatial_scale, + ) + + # One materialize: contiguous FHWC in latent.dtype, still [-1, 1]. + # CFHW→FHWC cannot be inplace; range mapping is left to to_rgb. + cfhw = _narrow_content_cfhw(buf, frames_keep) + src = cfhw[0] # C, F, H, W (view into accumulator) + video = torch.empty( + src.shape[1], + src.shape[2], + src.shape[3], + src.shape[0], + dtype=latent.dtype, + device=src.device, + ) + video.copy_(src.permute(1, 2, 3, 0)) + if not complementary: + assert wts is not None + w_cfhw = _narrow_content_cfhw(wts, frames_keep) + wview = w_cfhw[0].permute(1, 2, 3, 0) + # Inplace floor on exclusive weight region only (discarded after emit). + wview.clamp_min_(diffusion_tiling._weight_floor(w_cfhw.dtype)) + video.div_(wview) + return video + + for gi, group in enumerate(groups): + curr_temporal_slice = group_slices[gi] + buffer, weights = self._decode_temporal_group_isolated( + group, + feat_s4, + s4_t, + x_t_init, + timestep, + full_video_shape, + curr_temporal_slice, + generator=generator, + complementary=complementary, + ) + + if overlap_stub is not None: + overlap_len = int(overlap_stub.shape[2]) + if overlap_len > 0: + # Stub is exactly the region overlapping this group (cloned when + # the previous group finished); blend then write back into buffer. + overlap_stub += buffer[:, :, :overlap_len] + if complementary: + buffer[:, :, :overlap_len] = overlap_stub + else: + assert overlap_stub_weights is not None + assert weights is not None + overlap_stub_weights += weights[:, :, :overlap_len] + buffer[:, :, :overlap_len] = overlap_stub + weights[:, :, :overlap_len] = overlap_stub_weights + overlap_stub = None + overlap_stub_weights = None + + if gi + 1 < len(groups): + next_start = group_slices[gi + 1].start + exclusive_len = min(max(0, next_start - curr_temporal_slice.start), buffer.shape[2]) + emitted = _crop_emit( + buffer[:, :, :exclusive_len], + None if weights is None else weights[:, :, :exclusive_len], + curr_temporal_slice.start, + ) + if emitted is not None: + yield emitted + # Retain only the trailing overlap for the next handoff. + overlap_stub = buffer[:, :, exclusive_len:].clone() + if not complementary: + assert weights is not None + overlap_stub_weights = weights[:, :, exclusive_len:].clone() + del buffer, weights + else: + emitted = _crop_emit(buffer, weights, curr_temporal_slice.start) + if emitted is not None: + yield emitted + + def forward( + self, + sample: torch.Tensor, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + """Decode via ``_decode_pixels`` with ``tiling_config=None`` (single full tile).""" + return next(self._decode_pixels(sample, tiling_config=None, generator=generator)) + + def recommended_tiling_config( + self, + *, + height: int, + width: int, + num_frames: int, + mode: DiffVAEMode = DiffVAEMode.CHUNKED_EAGER, + free_bytes: int | None = None, + model_bytes: int | None = None, + ) -> TileSizeConfig: + """DiffVAE-aware tiling: stage-4/5 overlaps + memory-capped tile sizes.""" + if free_bytes is None: + device = next(self.parameters()).device + free_bytes = cuda_activation_budget_bytes(device) if device.type == "cuda" else 0 + if model_bytes is None: + model_bytes = sum(p.numel() * p.element_size() for p in self.parameters()) + pixel_scale = diffusion_tiling.stage4_to_pixel_scale_factors( + tuple(self.upsamples[3].stride), # type: ignore[arg-type] + self.patch_size, + ) + return diffusion_tiling.recommended_decode_tiling_config( + tile_halos=self.tile_halos, + pixel_scale=pixel_scale, + min_tile_size_s4=self.tile_min_sizes, + patch_size=self.patch_size, + height=height, + width=width, + num_frames=num_frames, + mode=mode, + free_bytes=free_bytes, + stage5_channels=self.stage5_channels, + stage4_channels=self.stage_channels[3], + upsample_strides=tuple(tuple(upsample.stride) for upsample in self.upsamples), + model_bytes=model_bytes, + ) + + def tiled_decode( + self, + latent: torch.Tensor, + tiling_config: TilingConfig, + generator: torch.Generator | None = None, + ) -> Iterator[torch.Tensor]: + """Tiled decode: stages 1-3 once, stages 4-5 per tile, pixel blend.""" + yield from self._decode_pixels(latent, tiling_config, generator=generator) + + def decode_video( + self, + latent: torch.Tensor, + tiling_config: TilingConfig | None = None, + generator: torch.Generator | None = None, + ) -> Iterator[torch.Tensor]: + """Decode latent video, yielding float chunk(s) ``[f, h, w, c]`` in ``[0, 1]``. + Untiled and tiled both go through ``_decode_pixels``. Tiled decode may yield + multiple times when ``tiling_config.frames`` splits the video. + Layout is packed once to contiguous FHWC on emit; ``to_rgb`` only does + inplace ``[-1, 1]→[0, 1]`` (no second realloc). + """ + + def to_rgb(frames: torch.Tensor) -> torch.Tensor: + return frames.add_(1).mul_(0.5).clamp_(0, 1) + + for chunk in self._decode_pixels(latent, tiling_config, generator=generator, as_fhwc=True): + yield to_rgb(chunk) + + @classmethod + def from_checkpoint( + cls, + checkpoint_path: str | Path, + *, + device: torch.device | str = "cpu", + torch_dtype: torch.dtype = torch.bfloat16, + ) -> "DiffusionVideoDecoder": + """Construct and strictly load the standalone LTX-2.5 DiffVAE decoder.""" + checkpoint = inspect_checkpoint(checkpoint_path) + vae_config = checkpoint.config.get("vae") + if not isinstance(vae_config, dict) or not isinstance(vae_config.get("decoder"), dict): + raise ValueError("LTX-2.5 DiffVAE checkpoint is missing vae.decoder config") + decoder_config = vae_config["decoder"] + kwargs = { + "in_channels": decoder_config["in_channels"], + "out_channels": decoder_config["out_channels"], + "patch_size": decoder_config["patch_size"], + "head_dim": decoder_config["head_dim"], + "stage_channels": tuple(decoder_config["stage_channels"]), + "stage_depths": tuple(decoder_config["stage_depths"]), + "stage_kernels": tuple(tuple(value) for value in decoder_config["stage_kernels"]), + "upsamples": tuple((tuple(value[0]), value[1]) for value in decoder_config["upsamples"]), + "stage5_kernel": tuple(decoder_config["stage5_kernel"]), + "default_num_inference_steps": decoder_config["default_num_inference_steps"], + "timestep_scale_multiplier": decoder_config["timestep_scale_multiplier"], + "model_output_type": vae_config["model_output_type"], + } + with torch.device("meta"): + model = cls(**kwargs) + unexpected, missing = ltx25_diffusion_vae_checkpoint_key_coverage(checkpoint.path, set(model.state_dict())) + if unexpected or missing: + raise ValueError( + "LTX-2.5 DiffVAE checkpoint coverage mismatch: " + f"unexpected={sorted(unexpected)[:5]}, missing={sorted(missing)[:5]}" + ) + state_dict = _load_ltx25_diffusion_vae_state_dict(checkpoint.path) + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=True, assign=True) + if missing_keys or unexpected_keys: + raise ValueError( + f"LTX-2.5 DiffVAE load mismatch: missing={missing_keys[:5]}, unexpected={unexpected_keys[:5]}" + ) + model = model.to(device=device, dtype=torch_dtype) + for module in model.modules(): + if isinstance(module, NeighborhoodAttention3D): + # The official loader keeps these non-persistent frequency buffers + # in FP32 even when decoder weights run in BF16. + rope_dim_t, rope_dim_h, rope_dim_w = module.rope_dim_split + module.rope_inv_t = rope_inv_freqs(rope_dim_t, module.rope_base).to(device=module.rope_inv_t.device) + module.rope_inv_h = rope_inv_freqs(rope_dim_h, module.rope_base).to(device=module.rope_inv_h.device) + module.rope_inv_w = rope_inv_freqs(rope_dim_w, module.rope_base).to(device=module.rope_inv_w.device) + _configure_chunked_eager_mode(model) + return model.eval() + + +def _configure_chunked_eager_mode(model: DiffusionVideoDecoder) -> None: + """Apply upstream's default eager DiffVAE recipe to an isolated decoder.""" + model.deferred_stage4_upsample = True + for block in model.diff_blocks: + if isinstance(block, DiffusionNABlock): + block.__class__ = ChunkedDiffusionNABlock + block.stage4_upsample = model.upsamples[3] + configure_w_chunks(block, w_chunks=4) + + attention = NattenAttention("cutlass-fna") if natten_available() else fallback_na_attention() + for module in model.modules(): + if isinstance(module, NeighborhoodAttention3D): + module.attention_function = attention + module.natten_backend = "cutlass-fna" if natten_available() else None + + +def _configure_chunked_compile_mode(model: DiffusionVideoDecoder) -> None: + """Compile chunked residuals with the conservative official cutlass-fna NATTEN backend.""" + model.deferred_stage4_upsample = True + for block in model.diff_blocks: + if isinstance(block, DiffusionNABlock): + block.__class__ = ChunkedDiffusionNABlock + block.stage4_upsample = model.upsamples[3] + configure_w_chunks(block, w_chunks=4) + + for module in model.modules(): + if isinstance(module, NeighborhoodAttention3D): + module.attention_function = NattenAttention("cutlass-fna") + module.natten_backend = "cutlass-fna" + compile_diffusion_decoder(model) + + +def _diffusion_vae_targets(key: str) -> tuple[str, ...]: + if key.startswith("decoder."): + target = key.removeprefix("decoder.") + # The official checkpoint carries this vestigial tensor; NADiffusionDecoder + # has no corresponding parameter and upstream does not consume it. + if target == "type_emb": + return () + if target.endswith(".qkv.weight"): + prefix = target.removesuffix(".qkv.weight") + return tuple(f"{prefix}.qkv.to_{part}.weight" for part in ("q", "k", "v")) + if target.endswith(".qkv.bias"): + prefix = target.removesuffix(".qkv.bias") + return tuple(f"{prefix}.qkv.to_{part}.bias" for part in ("q", "k", "v")) + return ( + target.replace("t_embedder.mlp.0", "t_embedder.timestep_embedder.linear_1").replace( + "t_embedder.mlp.2", "t_embedder.timestep_embedder.linear_2" + ), + ) + if key.startswith("per_channel_statistics."): + return (key,) + return () + + +def _load_ltx25_diffusion_vae_state_dict(checkpoint_path: str | Path) -> dict[str, torch.Tensor]: + state_dict: dict[str, torch.Tensor] = {} + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as checkpoint: + for key in checkpoint.keys(): + targets = _diffusion_vae_targets(key) + if not targets: + continue + value = checkpoint.get_tensor(key) + if len(targets) == 3: + values = value.chunk(3, dim=0) + state_dict.update(dict(zip(targets, values, strict=True))) + else: + state_dict[targets[0]] = value + return state_dict + + +def ltx25_diffusion_vae_checkpoint_key_coverage( + checkpoint_path: str | Path, + model_keys: set[str], +) -> tuple[set[str], set[str]]: + """Return unexplained source keys and missing isolated DiffVAE keys.""" + mapped: set[str] = set() + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as checkpoint: + for key in checkpoint.keys(): + mapped.update(_diffusion_vae_targets(key)) + return mapped - model_keys, model_keys - mapped diff --git a/telefuser/models/ltx25/diff_vae/enums.py b/telefuser/models/ltx25/diff_vae/enums.py new file mode 100644 index 0000000..edb2eac --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/enums.py @@ -0,0 +1,20 @@ +from enum import Enum + + +class NormLayerType(Enum): + GROUP_NORM = "group_norm" + PIXEL_NORM = "pixel_norm" + + +class LogVarianceType(Enum): + PER_CHANNEL = "per_channel" + UNIFORM = "uniform" + CONSTANT = "constant" + NONE = "none" + + +class PaddingModeType(Enum): + ZEROS = "zeros" + REFLECT = "reflect" + REPLICATE = "replicate" + CIRCULAR = "circular" diff --git a/telefuser/models/ltx25/diff_vae/ops.py b/telefuser/models/ltx25/diff_vae/ops.py new file mode 100644 index 0000000..7eccea5 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/ops.py @@ -0,0 +1,84 @@ +import torch +from einops import rearrange +from torch import nn + + +def patchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: + """ + Rearrange spatial dimensions into channels. Divides image into patch_size x patch_size blocks + and moves pixels from each block into separate channels (space-to-depth). + Args: + x: Input tensor (4D or 5D) + patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, divides HxW into 4x4 blocks. + patch_size_t: Temporal patch size for frames. Default=1 (no temporal patching). + For 5D: (B, C, F, H, W) -> (B, Cx(patch_size_hw^2)x(patch_size_t), F/patch_size_t, H/patch_size_hw, W/patch_size_hw) + Example: (B, 3, 33, 512, 512) with patch_size_hw=4, patch_size_t=1 -> (B, 48, 33, 128, 128) + """ + if patch_size_hw == 1 and patch_size_t == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size_hw, r=patch_size_hw) + elif x.dim() == 5: + x = rearrange( + x, + "b c (f p) (h q) (w r) -> b (c p r q) f h w", + p=patch_size_t, + q=patch_size_hw, + r=patch_size_hw, + ) + else: + raise ValueError(f"Invalid input shape: {x.shape}") + + return x + + +def unpatchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: + """ + Rearrange channels back into spatial dimensions. Inverse of patchify - moves pixels from + channels back into patch_size x patch_size blocks (depth-to-space). + Args: + x: Input tensor (4D or 5D) + patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, expands HxW by 4x. + patch_size_t: Temporal patch size for frames. Default=1 (no temporal expansion). + For 5D: (B, Cx(patch_size_hw^2)x(patch_size_t), F, H, W) -> (B, C, Fxpatch_size_t, Hxpatch_size_hw, Wxpatch_size_hw) + Example: (B, 48, 33, 128, 128) with patch_size_hw=4, patch_size_t=1 -> (B, 3, 33, 512, 512) + """ + if patch_size_hw == 1 and patch_size_t == 1: + return x + + if x.dim() == 4: + x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size_hw, r=patch_size_hw) + elif x.dim() == 5: + x = rearrange( + x, + "b (c p r q) f h w -> b c (f p) (h q) (w r)", + p=patch_size_t, + q=patch_size_hw, + r=patch_size_hw, + ) + + return x + + +class PerChannelStatistics(nn.Module): + """ + Per-channel statistics for normalizing and denormalizing the latent representation. + This statics is computed over the entire dataset and stored in model's checkpoint under VAE state_dict. + Defaults are identity (std=1, mean=0) so models constructed without a checkpoint + do not inherit allocator garbage / NaNs from ``torch.empty``. + """ + + def __init__(self, latent_channels: int = 128): + super().__init__() + self.register_buffer("std-of-means", torch.ones(latent_channels)) + self.register_buffer("mean-of-means", torch.zeros(latent_channels)) + + def un_normalize(self, x: torch.Tensor) -> torch.Tensor: + return (x * self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)) + self.get_buffer("mean-of-means").view( + 1, -1, 1, 1, 1 + ).to(x) + + def normalize(self, x: torch.Tensor) -> torch.Tensor: + return (x - self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)) / self.get_buffer("std-of-means").view( + 1, -1, 1, 1, 1 + ).to(x) diff --git a/telefuser/models/ltx25/diff_vae/runtime.py b/telefuser/models/ltx25/diff_vae/runtime.py new file mode 100644 index 0000000..3858c46 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/runtime.py @@ -0,0 +1,92 @@ +"""Runtime primitives required by the isolated LTX-2.5 DiffVAE path.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Callable, Protocol + +import torch +from torch import nn + + +class AttentionCallable(Protocol): + """Callable interface used by the ConvVAE spatial-attention blocks.""" + + def __call__(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int) -> torch.Tensor: ... + + +def _sdpa_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int) -> torch.Tensor: + batch, _, channels = q.shape + head_dim = channels // heads + q, k, v = (tensor.view(batch, -1, heads, head_dim).transpose(1, 2) for tensor in (q, k, v)) + output = torch.nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False) + return output.transpose(1, 2).reshape(batch, -1, channels) + + +class AttentionFunction(Enum): + """Supported eager attention backend for the isolated ConvVAE.""" + + PYTORCH = "pytorch" + + def to_callable(self) -> AttentionCallable: + """Resolve the configured backend without importing upstream LTX-Core.""" + return _sdpa_attention + + +class PixelNorm(nn.Module): + """Per-location RMS normalization used by the LTX video VAE.""" + + def __init__(self, dim: int = 1, eps: float = 1e-8) -> None: + super().__init__() + self.dim = dim + self.eps = eps + + def forward(self, value: torch.Tensor) -> torch.Tensor: + """Normalize values over the configured channel dimension.""" + return value / torch.sqrt(torch.mean(value.square(), dim=self.dim, keepdim=True) + self.eps) + + +class Disposable: + """Mixin that releases parameters and persistent buffers to the meta device.""" + + def dispose(self) -> None: + """Release tensor storage while preserving the module structure.""" + if not isinstance(self, nn.Module): + raise TypeError(f"{type(self).__name__} must be an nn.Module to dispose") + persistent = set(self.state_dict()) + for name, parameter in list(self.named_parameters()): + parent_name, _, attribute = name.rpartition(".") + parent = self.get_submodule(parent_name) if parent_name else self + setattr( + parent, + attribute, + nn.Parameter(torch.empty_like(parameter, device="meta"), requires_grad=parameter.requires_grad), + ) + for name, buffer in list(self.named_buffers()): + if name not in persistent: + continue + parent_name, _, attribute = name.rpartition(".") + parent = self.get_submodule(parent_name) if parent_name else self + parent.register_buffer(attribute, torch.empty_like(buffer, device="meta"), persistent=True) + + +@dataclass(frozen=True) +class CompilationConfig: + """Minimal compile settings accepted by the isolated DiffVAE implementation.""" + + mode: str | None = None + backend: str | None = None + fullgraph: bool = False + dynamic: bool = False + inductor_config: dict[str, Any] = field(default_factory=dict) + dynamo_config: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class ModuleOps: + """Local model mutator used by DiffVAE configuration helpers.""" + + name: str + matcher: Callable[[nn.Module], bool] + mutator: Callable[[nn.Module], nn.Module] diff --git a/telefuser/models/ltx25/diff_vae/tiling.py b/telefuser/models/ltx25/diff_vae/tiling.py new file mode 100644 index 0000000..e6387c2 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/tiling.py @@ -0,0 +1,945 @@ +from __future__ import annotations + +import itertools +import math +from dataclasses import dataclass, replace +from typing import Callable, NamedTuple, Sequence + +import torch + +from .types import SpatioTemporalScaleFactors, VideoPixelShape + + +def compute_trapezoidal_mask_1d( + length: int, + ramp_left: int, + ramp_right: int, + left_starts_from_0: bool = False, +) -> torch.Tensor: + """ + Generate a 1D trapezoidal blending mask with linear ramps. + Args: + length: Output length of the mask. + ramp_left: Fade-in length on the left. + ramp_right: Fade-out length on the right. + left_starts_from_0: Whether the ramp starts from 0 or first non-zero value. + Useful for temporal tiles where the first tile is causal. + Returns: + A 1D tensor of shape `(length,)` with values in [0, 1]. + """ + if length <= 0: + raise ValueError("Mask length must be positive.") + + ramp_left = max(0, min(ramp_left, length)) + ramp_right = max(0, min(ramp_right, length)) + + mask = torch.ones(length) + + if ramp_left > 0: + interval_length = ramp_left + 1 if left_starts_from_0 else ramp_left + 2 + fade_in = torch.linspace(0.0, 1.0, interval_length)[:-1] + if not left_starts_from_0: + fade_in = fade_in[1:] + mask[:ramp_left] *= fade_in + + if ramp_right > 0: + fade_out = torch.linspace(1.0, 0.0, steps=ramp_right + 2)[1:-1] + mask[-ramp_right:] *= fade_out + + return mask.clamp_(0, 1) + + +def compute_rectangular_mask_1d( + length: int, + left_ramp: int, + right_ramp: int, +) -> torch.Tensor: + """ + Generate a 1D rectangular (pulse) mask. + Args: + length: Output length of the mask. + left_ramp: Number of elements at the start of the mask to set to 0. + right_ramp: Number of elements at the end of the mask to set to 0. + Returns: + A 1D tensor of shape `(length,)` with values 0 or 1. + """ + if length <= 0: + raise ValueError("Mask length must be positive.") + + mask = torch.ones(length) + if left_ramp > 0: + mask[:left_ramp] = 0 + if right_ramp > 0: + mask[-right_ramp:] = 0 + return mask + + +@dataclass(frozen=True) +class DimensionInterval: + start: int + end: int + left_ramp: int + right_ramp: int + + +@dataclass(frozen=True) +class DimensionIntervals: + """Intervals which a single dimension of the latent space is split into. + Each interval is defined by its start, end, left ramp, and right ramp. + The start and end are the indices of the first and last element (exclusive) in the interval. + Ramps are regions of the interval where the value of the mask tensor is + interpolated between 0 and 1 for blending with neighboring intervals. + The left ramp and right ramp values are the lengths of the left and right ramps. + """ + + intervals: list[DimensionInterval] + + +@dataclass(frozen=True) +class LatentIntervals: + """Intervals which the latent tensor of given shape is split into. + Each dimension of the latent space is split into intervals based on the length along said dimension. + """ + + original_shape: torch.Size + dimension_intervals: tuple[DimensionIntervals, ...] + + +# Operation to split a single dimension of the tensor into intervals based on the length along the dimension. +SplitOperation = Callable[[int], DimensionIntervals] +# Operation to map the intervals in input dimension to slices and masks along a corresponding output dimension. +MappingOperation = Callable[[DimensionIntervals], tuple[list[slice], list[torch.Tensor]]] + + +def default_split_operation(length: int) -> DimensionIntervals: + return DimensionIntervals(intervals=[DimensionInterval(start=0, end=length, left_ramp=0, right_ramp=0)]) + + +DEFAULT_SPLIT_OPERATION: SplitOperation = default_split_operation + + +def untiled_mask_1d() -> torch.Tensor: + """Length-1 ones that broadcast over an untiled axis (historical ``None`` mask).""" + return torch.ones(1) + + +def default_mapping_operation( + _intervals: DimensionIntervals, +) -> tuple[list[slice], list[torch.Tensor]]: + return [slice(0, None)], [untiled_mask_1d()] + + +DEFAULT_MAPPING_OPERATION: MappingOperation = default_mapping_operation + + +# --------------------------------------------------------------------------- +# Split functions +# --------------------------------------------------------------------------- + + +def _grow_last_tile_to_min(intervals: list[DimensionInterval], min_tile_size: int) -> list[DimensionInterval]: + """Grow a short last tile left to ``min_tile_size``; widen penultimate ``right_ramp``.""" + if len(intervals) <= 1: + return list(intervals) + last = intervals[-1] + if last.end - last.start >= min_tile_size: + return list(intervals) + new_start = last.end - min_tile_size + prev = intervals[-2] + new_overlap = prev.end - new_start + return [ + *intervals[:-2], + replace(prev, right_ramp=new_overlap), + replace(last, start=new_start, left_ramp=new_overlap), + ] + + +def _validate_tile_intervals(intervals: list[DimensionInterval], *, dim_size: int, min_tile_size: int) -> None: + """Validate coverage, ramp/overlap consistency, and ``min_tile_size``.""" + if not intervals or intervals[0].start != 0 or intervals[-1].end != dim_size: + raise ValueError(f"tiles must cover [0, {dim_size})") + for i, iv in enumerate(intervals): + length = iv.end - iv.start + if length < min_tile_size: + raise ValueError(f"tile {i} length {length} is below min_tile_size={min_tile_size}") + if iv.left_ramp < 0 or iv.right_ramp < 0 or iv.left_ramp > length or iv.right_ramp > length: + raise ValueError(f"tile {i} has invalid ramps: left={iv.left_ramp}, right={iv.right_ramp}, length={length}") + if i == 0: + continue + overlap = intervals[i - 1].end - iv.start + if overlap < 0 or intervals[i - 1].right_ramp != overlap or iv.left_ramp != overlap: + raise ValueError(f"tiles {i - 1}/{i}: ramp/overlap mismatch (overlap={overlap})") + + +def split_by_size(size: int, overlap: int, min_tile_size: int | None = None) -> SplitOperation: + """Split a dimension into overlapping tiles of a given size. + Tiles are sized ``size`` with ``overlap`` shared elements between + consecutive tiles. The last tile may be shorter if the dimension + doesn't divide evenly. If ``min_tile_size`` is set and the last tile is + shorter, it is grown leftward (penultimate ``right_ramp`` widens); the + result is validated and invalid layouts raise ``ValueError``. + Args: + size: Target tile size (in axis units). + overlap: Overlap between consecutive tiles. + min_tile_size: Optional minimum tile length. ``None`` keeps legacy + short-last-tile behavior. + Returns: + A split operation that divides a dimension into tiles. + """ + if size <= 0: + raise ValueError(f"size must be > 0, got {size}") + if overlap < 0 or overlap >= size: + raise ValueError(f"overlap must satisfy 0 <= overlap < size, got overlap={overlap}, size={size}") + if min_tile_size is not None and min_tile_size < 1: + raise ValueError(f"min_tile_size must be >= 1, got {min_tile_size}") + + def split(dimension_size: int) -> DimensionIntervals: + if min_tile_size is not None and dimension_size < min_tile_size: + return DEFAULT_SPLIT_OPERATION(dimension_size) + if dimension_size <= size: + return DEFAULT_SPLIT_OPERATION(dimension_size) + amount = (dimension_size + size - 2 * overlap - 1) // (size - overlap) + intervals = [ + DimensionInterval(start=0, end=size, left_ramp=0, right_ramp=overlap), + *( + DimensionInterval( + start=i * (size - overlap), + end=i * (size - overlap) + size, + left_ramp=overlap, + right_ramp=overlap, + ) + for i in range(1, amount - 1) + ), + DimensionInterval( + start=(amount - 1) * (size - overlap), end=dimension_size, left_ramp=overlap, right_ramp=0 + ), + ] + if min_tile_size is not None: + intervals = _grow_last_tile_to_min(intervals, min_tile_size) + _validate_tile_intervals(intervals, dim_size=dimension_size, min_tile_size=min_tile_size) + return DimensionIntervals(intervals=intervals) + + return split + + +def split_temporal_causal(size: int, overlap: int, min_tile_size: int | None = None) -> SplitOperation: + """Split a temporal axis into overlapping tiles with causal handling. + Each tile after the first is shifted back by 1 and its left ramp is + increased by 1, ensuring causal continuity through the blend ramps. + Args: + size: Tile size in axis units. + overlap: Overlap between tiles in the same units. + min_tile_size: Optional floor forwarded to :func:`split_by_size`. + Returns: + Split operation that divides temporal dimension with causal handling. + """ + non_causal_split = split_by_size(size, overlap, min_tile_size=min_tile_size) + + def split(dimension_size: int) -> DimensionIntervals: + if dimension_size <= size: + return DEFAULT_SPLIT_OPERATION(dimension_size) + dim_intervals = non_causal_split(dimension_size) + if len(dim_intervals.intervals) <= 1: + return dim_intervals + modified_intervals = [dim_intervals.intervals[0]] + [ + replace(interval, start=interval.start - 1, left_ramp=interval.left_ramp + 1) + for interval in dim_intervals.intervals[1:] + ] + return DimensionIntervals(intervals=modified_intervals) + + return split + + +def split_temporal(tile_size_frames: int, overlap_frames: int) -> SplitOperation: + """Split a temporal axis in video frame space into overlapping tiles. + Args: + tile_size_frames: Tile length in frames. + overlap_frames: Overlap between consecutive tiles in frames. + Returns: + Split operation that takes frame count and returns DimensionIntervals in frame indices. + """ + non_causal_split = split_by_size(tile_size_frames, overlap_frames) + + def split(dimension_size: int) -> DimensionIntervals: + if dimension_size <= tile_size_frames: + return DEFAULT_SPLIT_OPERATION(dimension_size) + dim_intervals = non_causal_split(dimension_size) + modified_intervals = [ + replace(interval, end=interval.end + 1, right_ramp=0) for interval in dim_intervals.intervals[:-1] + ] + [replace(dim_intervals.intervals[-1], right_ramp=0)] + return DimensionIntervals(intervals=modified_intervals) + + return split + + +def split_by_count_temporal_causal( + num_tiles: int, overlap: int = 0, min_tile_size: int | None = None +) -> SplitOperation: + """Split a temporal dimension by count with causal handling. + Wraps :func:`split_by_count` with the same causal adjustment as + :func:`split_temporal_causal`: each tile after the first is shifted + back by 1 and its left ramp is increased by 1. + Args: + num_tiles: Number of tiles. Must be >= 1. + overlap: Overlap between adjacent tiles (default 0). + min_tile_size: Optional floor forwarded to :func:`split_by_count`. + Returns: + A split operation that divides a temporal dimension into tiles. + """ + non_causal_split = split_by_count(num_tiles, overlap, min_tile_size=min_tile_size) + + def split(dimension_size: int) -> DimensionIntervals: + dim_intervals = non_causal_split(dimension_size) + if len(dim_intervals.intervals) <= 1: + return dim_intervals + modified_intervals = [dim_intervals.intervals[0]] + [ + replace(interval, start=interval.start - 1, left_ramp=interval.left_ramp + 1) + for interval in dim_intervals.intervals[1:] + ] + return DimensionIntervals(intervals=modified_intervals) + + return split + + +def split_by_count(num_tiles: int, overlap: int = 0, min_tile_size: int | None = None) -> SplitOperation: + """Split a dimension into a given number of tiles with overlap. + Computes the tile size as + ``(dim_size + overlap * (num_tiles - 1)) // num_tiles`` so that + ``num_tiles`` tiles of that size with ``overlap`` shared elements + cover the dimension evenly. Delegates to :func:`split_by_size` for + the actual interval construction. + When the total ``dim_size + overlap * (num_tiles - 1)`` is not evenly + divisible by ``num_tiles``, the first ``remainder`` tiles each absorb + one extra unit. + Args: + num_tiles: Number of tiles. Must be >= 1. + overlap: Overlap between adjacent tiles (default 0). Must be >= 0 + and less than the computed tile size. + min_tile_size: Optional floor forwarded to last-tile growth / validation. + Returns: + A split operation that divides a dimension into tiles. + """ + if num_tiles < 1: + raise ValueError(f"num_tiles must be >= 1, got {num_tiles}") + if overlap < 0: + raise ValueError(f"overlap must be >= 0, got {overlap}") + if min_tile_size is not None and min_tile_size < 1: + raise ValueError(f"min_tile_size must be >= 1, got {min_tile_size}") + + def split(dim_size: int) -> DimensionIntervals: + if num_tiles > dim_size: + raise ValueError( + f"num_tiles ({num_tiles}) exceeds dim_size ({dim_size}). Cannot assign at least 1 unit per tile." + ) + if num_tiles == 1: + return DEFAULT_SPLIT_OPERATION(dim_size) + + total = dim_size + overlap * (num_tiles - 1) + tile_size = total // num_tiles + if tile_size <= overlap: + raise ValueError( + f"split_by_count produced size={tile_size} <= overlap={overlap} " + f"for dim_size={dim_size}, num_tiles={num_tiles}" + ) + remainder = total % num_tiles + + base_intervals = split_by_size(tile_size, overlap)(dim_size - remainder).intervals + + # First `remainder` tiles each absorb 1 extra unit; shift subsequent boundaries. + intervals: list[DimensionInterval] = [] + for i, iv in enumerate(base_intervals): + shift = min(i, remainder) + grow = 1 if i < remainder else 0 + intervals.append(replace(iv, start=iv.start + shift, end=iv.end + shift + grow)) + + if min_tile_size is not None: + intervals = _grow_last_tile_to_min(intervals, min_tile_size) + _validate_tile_intervals(intervals, dim_size=dim_size, min_tile_size=min_tile_size) + + return DimensionIntervals(intervals=intervals) + + return split + + +# --------------------------------------------------------------------------- +# Mapping operations +# --------------------------------------------------------------------------- + + +def identity_mapping_operation(intervals: DimensionIntervals) -> tuple[list[slice], list[torch.Tensor]]: + """Map each DimensionInterval to an output region at the same position, with trapezoidal blend masks. + For every interval the output start/end matches the input start/end and a + 1-D blending mask is built from the interval's left_ramp and right_ramp. + """ + out_slices: list[slice] = [] + masks: list[torch.Tensor] = [] + for iv in intervals.intervals: + out_slices.append(slice(iv.start, iv.end)) + masks.append(compute_trapezoidal_mask_1d(iv.end - iv.start, iv.left_ramp, iv.right_ramp)) + return out_slices, masks + + +class Tile(NamedTuple): + """ + Represents a single tile. + Attributes: + in_coords: + Tuple of slices specifying where to cut the tile from the INPUT tensor. + out_coords: + Tuple of slices specifying where this tile's OUTPUT should be placed in the reconstructed OUTPUT tensor. + masks_1d: + Per-dimension masks in OUTPUT units. + Untiled axes use a length-1 ones tensor (broadcasts). These are used + for separable blending (and for the dense ``blend_mask`` property). + Methods: + blend_mask: + Create a single N-D mask from the per-dimension masks. + """ + + in_coords: tuple[slice, ...] + out_coords: tuple[slice, ...] + masks_1d: tuple[torch.Tensor, ...] + + @property + def blend_mask(self) -> torch.Tensor: + num_dims = len(self.out_coords) + per_dimension_masks: list[torch.Tensor] = [] + + for dim_idx in range(num_dims): + mask_1d = self.masks_1d[dim_idx] + view_shape = [1] * num_dims + # Reshape (L,) -> (1, ..., L, ..., 1) so masks across dimensions broadcast-multiply. + view_shape[dim_idx] = mask_1d.shape[0] + per_dimension_masks.append(mask_1d.view(*view_shape)) + + # Multiply per-dimension masks to form the full N-D mask (separable blending window). + combined_mask = per_dimension_masks[0] + for mask in per_dimension_masks[1:]: + combined_mask = combined_mask * mask + + return combined_mask + + +def scale_by_masks_1d(x: torch.Tensor, masks_1d: Sequence[torch.Tensor]) -> torch.Tensor: + """Multiply ``x`` by separable 1d masks with broadcasting. + ``len(masks_1d)`` must equal ``x.ndim``. Prefer float32 masks so bf16/fp16 ``x`` promotes. + Length-1 masks (untiled axes) broadcast over that dimension. + """ + if len(masks_1d) != x.ndim: + raise ValueError(f"masks_1d length {len(masks_1d)} != x.ndim {x.ndim}") + out = x + for axis, mask in enumerate(masks_1d): + view_shape = [1] * x.ndim + view_shape[axis] = -1 + out = out * mask.reshape(*view_shape) + return out + + +def masks_are_complementary( + tiles: Sequence[Tile], + full_shape: Sequence[int], + *, + atol: float = 1e-5, +) -> bool: + """Return whether per-axis 1d blend masks partition unity (sum to 1). + Checks each axis independently over the unique out-slices on that axis + (cartesian tile products would otherwise multi-count the same 1d interval). + When True, weighted accumulation needs no denominator. + """ + if not tiles: + return True + ndim = len(full_shape) + for tile in tiles: + if len(tile.out_coords) != ndim or len(tile.masks_1d) != ndim: + raise ValueError( + f"Tile out_coords/masks_1d rank {len(tile.out_coords)}/{len(tile.masks_1d)} != full_shape rank {ndim}" + ) + for axis, length in enumerate(full_shape): + # Explicit CPU float32: masks may live on CUDA; a non-CPU default device + # must not place ``acc`` on GPU (device-mismatch on ``acc[sl] +=``). + acc = torch.zeros(length, dtype=torch.float32, device="cpu") + seen: set[tuple[int | None, int | None]] = set() + for tile in tiles: + sl = tile.out_coords[axis] + key = (sl.start, sl.stop) + if key in seen: + continue + seen.add(key) + # Length-1 untiled masks broadcast over ``acc[sl]``. + acc[sl] += tile.masks_1d[axis].detach().float().cpu() + if not torch.allclose(acc, torch.ones(length, dtype=torch.float32), atol=atol, rtol=0.0): + return False + return True + + +def compute_summed_weights( + tiles: Sequence[Tile], + full_shape: Sequence[int], +) -> torch.Tensor: + """Build the dense denominator for weighted blending over ``full_shape``. + Uses separable per-axis mask broadcasts — never ``Tile.blend_mask``. + Requires concrete ``out_coords`` (``stop`` not ``None``) on every axis. + Always builds on CPU float32 so CUDA masks / a non-CPU default device cannot + place a multi-GB ``[F,H,W]`` tensor on GPU. + """ + weights = torch.zeros(*full_shape, dtype=torch.float32, device="cpu") + for tile in tiles: + masks = tuple(m.detach().float().cpu() for m in tile.masks_1d) + region_shape = tuple(s.stop - s.start for s in tile.out_coords) + region = torch.ones(region_shape, dtype=torch.float32, device="cpu") + weights[tile.out_coords] += scale_by_masks_1d(region, masks) + return weights.clamp(min=1e-8) + + +def create_tiles_from_intervals_and_mappers( + intervals: LatentIntervals, + mappers: list[MappingOperation], +) -> list[Tile]: + full_dim_input_slices: list[list[slice]] = [] + full_dim_output_slices: list[list[slice]] = [] + full_dim_masks_1d: list[list[torch.Tensor]] = [] + for axis_index in range(len(intervals.original_shape)): + dimension_intervals = intervals.dimension_intervals[axis_index] + input_slices = [slice(interval.start, interval.end) for interval in dimension_intervals.intervals] + output_slices, masks_1d = mappers[axis_index](dimension_intervals) + n_intervals = len(input_slices) + if len(output_slices) != n_intervals or len(masks_1d) != n_intervals: + raise ValueError( + f"Axis {axis_index}: mapper produced {len(output_slices)} output slices and " + f"{len(masks_1d)} masks for {n_intervals} input intervals" + ) + full_dim_input_slices.append(input_slices) + full_dim_output_slices.append(output_slices) + full_dim_masks_1d.append(masks_1d) + + return [ + Tile(in_coords=in_coord, out_coords=out_coord, masks_1d=mask_1d) + for in_coord, out_coord, mask_1d in zip( + itertools.product(*full_dim_input_slices), + itertools.product(*full_dim_output_slices), + itertools.product(*full_dim_masks_1d), + strict=True, + ) + ] + + +def create_tiles( + latent_shape: torch.Size, + splitters: list[SplitOperation], + mappers: list[MappingOperation], +) -> list[Tile]: + if len(splitters) != len(latent_shape): + raise ValueError( + f"Number of splitters must be equal to number of dimensions in latent shape, " + f"got {len(splitters)} and {len(latent_shape)}" + ) + if len(mappers) != len(latent_shape): + raise ValueError( + f"Number of mappers must be equal to number of dimensions in latent shape, " + f"got {len(mappers)} and {len(latent_shape)}" + ) + intervals = [splitter(length) for splitter, length in zip(splitters, latent_shape, strict=True)] + latent_intervals = LatentIntervals(original_shape=latent_shape, dimension_intervals=tuple(intervals)) + return create_tiles_from_intervals_and_mappers(latent_intervals, mappers) + + +def group_tiles_by_temporal_slice(tiles: list[Tile]) -> list[list[Tile]]: + """Group consecutive tiles that share the same temporal ``out_coords`` slice. + Assumes ``tiles`` is ordered with the temporal axis varying slowest (true + for every tile list this codebase builds via ``itertools.product`` with + the temporal axis first), so equal temporal slices are always contiguous. + """ + if not tiles: + return [] + + groups = [] + current_slice = tiles[0].out_coords[2] + current_group = [] + + for tile in tiles: + tile_slice = tile.out_coords[2] + if tile_slice == current_slice: + current_group.append(tile) + else: + groups.append(current_group) + current_slice = tile_slice + current_group = [tile] + + if current_group: + groups.append(current_group) + + return groups + + +# --------------------------------------------------------------------------- +# Video-grid tiling configs +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class DimensionTilingConfig: + """Tiling parameters for a single dimension of the patchified grid. + Attributes: + num_tiles: Number of tiles along this dimension. ``1`` with ``overlap=0`` + means the axis is not tiled. + overlap: Overlap between adjacent tiles, in latent grid units. + Adjacent tiles share ``overlap`` grid cells at their + boundary, producing an overlap zone blended with + trapezoidal masks. + """ + + num_tiles: int = 1 + overlap: int = 0 + + def __post_init__(self) -> None: + if self.num_tiles < 1: + raise ValueError(f"num_tiles must be >= 1, got {self.num_tiles}") + if self.overlap < 0: + raise ValueError(f"overlap must be >= 0, got {self.overlap}") + + def is_tiled(self) -> bool: + """True when this axis is split into more than one tile (or has overlap).""" + return self.num_tiles > 1 or self.overlap > 0 + + @classmethod + def from_tile_size(cls, dim_size: int, tile_size: int, overlap: int = 0) -> DimensionTilingConfig: + """Create config by computing ``num_tiles`` from dimension size and tile size. + Args: + dim_size: Total length of the dimension. + tile_size: Desired tile size. + overlap: Overlap between consecutive tiles. + Returns: + A ``DimensionTilingConfig`` with the computed ``num_tiles``. + """ + split_op = split_by_size(tile_size, overlap) + intervals = split_op(dim_size) + return cls(num_tiles=len(intervals.intervals), overlap=overlap) + + +@dataclass(frozen=True) +class DimensionSizeConfig: + """Tile size and overlap for a single video axis (frames / height / width). + Mirrors :class:`DimensionTilingConfig`, but specifies tile *size* rather than + tile *count*. ``tile_size=0`` means the axis is not tiled (covers the whole + length). Axis-specific VAE pixel constraints (divisibility / minimums) are + enforced by :meth:`TileSizeConfig.validate` for tiled axes only. + """ + + tile_size: int = 0 + overlap: int = 0 + + def __post_init__(self) -> None: + if self.tile_size < 0: + raise ValueError(f"tile_size must be >= 0, got {self.tile_size}") + if self.overlap < 0: + raise ValueError(f"overlap must be >= 0, got {self.overlap}") + if self.tile_size == 0: + if self.overlap != 0: + raise ValueError("untiled axis (tile_size=0) must have overlap=0") + return + if self.overlap >= self.tile_size: + raise ValueError(f"Overlap must be less than tile size, got {self.overlap} and {self.tile_size}") + + def is_tiled(self) -> bool: + """True when this axis has a positive tile size (caller intends to split it).""" + return self.tile_size > 0 + + +@dataclass(frozen=True) +class TileCountConfig: + """Tiling layout for a ``(F, H, W)`` grid by tile *counts*. + Overlaps are in latent-grid units. Mirror of :class:`TileSizeConfig`. + Attributes: + frames: Tiling along the temporal (frames) dimension. + height: Tiling along the latent height dimension. + width: Tiling along the latent width dimension. + """ + + frames: DimensionTilingConfig = DimensionTilingConfig() + height: DimensionTilingConfig = DimensionTilingConfig() + width: DimensionTilingConfig = DimensionTilingConfig() + + def validate(self, scale_factors: SpatioTemporalScaleFactors, video_shape: VideoPixelShape) -> None: + """Raise if this count layout cannot tile ``video_shape`` under ``scale_factors``. + Counts/overlaps are in latent-grid units. ``video_shape.frames <= 0`` skips the + temporal axis (duration not yet known). Spatial axes always checked. + """ + check_temporal = _assert_video_on_vae_grid(scale_factors, video_shape) + latent_h = video_shape.height // scale_factors.height + latent_w = video_shape.width // scale_factors.width + _validate_count_axis(self.height, latent_h, "height") + _validate_count_axis(self.width, latent_w, "width") + if check_temporal: + latent_f = (video_shape.frames - 1) // scale_factors.time + 1 + _validate_count_axis(self.frames, latent_f, "frames") + + def to_splitters( + self, + scale_factors: SpatioTemporalScaleFactors, + min_tile_size: tuple[int, int, int] | None = None, + *, + causal_temporal: bool = True, + ) -> tuple[SplitOperation, SplitOperation, SplitOperation]: + """Build ``(T, H, W)`` latent-grid split operations for this count layout. + ``scale_factors`` is accepted for signature parity with + :meth:`TileSizeConfig.to_splitters` and ignored — counts are already in + grid units. When ``causal_temporal`` is True (VAE encode/decode), the + frames axis uses :func:`split_by_count_temporal_causal`; otherwise plain + :func:`split_by_count`. ``min_tile_size`` is a per-axis floor in the same + units as the split. + """ + del scale_factors + min_t = min_h = min_w = None + if min_tile_size is not None: + min_t, min_h, min_w = min_tile_size + + def axis_split(cfg: DimensionTilingConfig, axis_min: int | None, *, temporal: bool) -> SplitOperation: + if not cfg.is_tiled(): + return DEFAULT_SPLIT_OPERATION + if temporal and causal_temporal: + return split_by_count_temporal_causal(cfg.num_tiles, cfg.overlap, min_tile_size=axis_min) + return split_by_count(cfg.num_tiles, cfg.overlap, min_tile_size=axis_min) + + return ( + axis_split(self.frames, min_t, temporal=True), + axis_split(self.height, min_h, temporal=False), + axis_split(self.width, min_w, temporal=False), + ) + + def video_chunks_number(self, num_frames: int) -> int: + """Number of temporal decode chunks for ``num_frames`` under this layout.""" + del num_frames + return max(1, self.frames.num_tiles) + + +@dataclass(frozen=True) +class TileSizeConfig: + """Size-based tiling layout for a ``(F, H, W)`` video — mirror of ``TileCountConfig``. + Each axis is a non-optional :class:`DimensionSizeConfig`; ``tile_size=0`` means + untiled on that axis (:meth:`DimensionSizeConfig.is_tiled`). Sizes and overlaps + are in pixel / frame units. Conversion to a split grid is an explicit + ``scale_factors`` argument to :meth:`to_splitters` (not stored on the config). + Legality vs a VAE grid is checked by :meth:`validate` (same factors decode will + pass to :meth:`to_splitters`), not at construction. + Attributes: + frames: Temporal tile size/overlap in frames. + height: Spatial height tile size/overlap in pixels. + width: Spatial width tile size/overlap in pixels. + """ + + frames: DimensionSizeConfig = DimensionSizeConfig() + height: DimensionSizeConfig = DimensionSizeConfig() + width: DimensionSizeConfig = DimensionSizeConfig() + + def validate(self, scale_factors: SpatioTemporalScaleFactors, video_shape: VideoPixelShape) -> None: + """Raise if this size layout is illegal for ``video_shape`` under ``scale_factors``. + Checks tile/overlap divisibility and minimums against the VAE grid, and that + the video extents are compatible with that grid. ``video_shape.frames <= 0`` + skips the temporal axis (duration not yet known); height/width always checked. + """ + check_temporal = _assert_video_on_vae_grid(scale_factors, video_shape) + _validate_size_axis(self.height, scale_factors.height, "height") + _validate_size_axis(self.width, scale_factors.width, "width") + if check_temporal: + _validate_size_axis(self.frames, scale_factors.time, "frames") + + @classmethod + def default(cls) -> TileSizeConfig: + return cls( + frames=DimensionSizeConfig(tile_size=80, overlap=24), + height=DimensionSizeConfig(tile_size=768, overlap=64), + width=DimensionSizeConfig(tile_size=768, overlap=64), + ) + + @classmethod + def from_long_side( + cls, + *, + long_side: DimensionSizeConfig, + height: int, + width: int, + scale_factors: SpatioTemporalScaleFactors, + frames: DimensionSizeConfig | None = None, + ) -> TileSizeConfig: + """Aspect-coupled construction — old single-spatial long-side behavior, explicit. + Matches main-era ``latent_tile_splitters``: scale the long-side tile in + *latent* units with ``round(size_lat * axis_lat / long_lat)``, then + multiply back by the VAE factor. Pixel-space ``round`` + ceil-snap would + bias the short axis up by almost one latent (e.g. 680 → 704 vs 672). + Both axes share ``long_side.overlap``. + """ + if height < 1 or width < 1: + raise ValueError(f"height/width must be >= 1, got {height}x{width}") + if not long_side.is_tiled(): + raise ValueError("long_side must be tiled (tile_size > 0)") + if scale_factors.height < 1 or scale_factors.width < 1: + raise ValueError(f"scale_factors height/width must be >= 1, got {scale_factors}") + span = max(height, width) + + def axis_size(axis_len: int, factor: int) -> int: + # Latent-grid round (same as main decode enable_on_axis), not pixel ceil. + axis_lat = axis_len // factor + long_lat = span // factor + size_lat = long_side.tile_size // factor + overlap_lat = long_side.overlap // factor + lower_threshold = max(2, overlap_lat + 1) + tile_lat = max(lower_threshold, round(size_lat * axis_lat / long_lat)) + tile_px = tile_lat * factor + min_legal = max(2 * factor, long_side.overlap + factor) + return max(tile_px, min_legal) + + return cls( + frames=DimensionSizeConfig() if frames is None else frames, + height=DimensionSizeConfig(tile_size=axis_size(height, scale_factors.height), overlap=long_side.overlap), + width=DimensionSizeConfig(tile_size=axis_size(width, scale_factors.width), overlap=long_side.overlap), + ) + + def to_splitters( + self, + scale_factors: SpatioTemporalScaleFactors, + min_tile_size: tuple[int, int, int] | None = None, + *, + causal_temporal: bool = True, + ) -> tuple[SplitOperation, SplitOperation, SplitOperation]: + """Build ``(T, H, W)`` grid split ops from pixel/frame sizes via ``scale_factors``. + When ``causal_temporal`` is True, frames use :func:`split_temporal_causal`. + """ + min_t = min_h = min_w = None + if min_tile_size is not None: + min_t, min_h, min_w = min_tile_size + + def enable_size_axis( + factor: int, + axis_min: int | None, + cfg: DimensionSizeConfig, + axis_name: str, + *, + temporal: bool, + ) -> SplitOperation: + if not cfg.is_tiled(): + return DEFAULT_SPLIT_OPERATION + _validate_size_axis(cfg, factor, axis_name) + size = cfg.tile_size // factor + overlap = cfg.overlap // factor + lower_threshold = max(2, overlap + 1) + tile = max(lower_threshold, size) + if temporal and causal_temporal: + return split_temporal_causal(tile, overlap, min_tile_size=axis_min) + return split_by_size(tile, overlap, min_tile_size=axis_min) + + return ( + enable_size_axis(scale_factors.time, min_t, self.frames, "frames", temporal=True), + enable_size_axis(scale_factors.height, min_h, self.height, "height", temporal=False), + enable_size_axis(scale_factors.width, min_w, self.width, "width", temporal=False), + ) + + def video_chunks_number(self, num_frames: int) -> int: + """Number of temporal decode chunks for ``num_frames`` under this layout.""" + if not self.frames.is_tiled(): + return 1 + frame_stride = self.frames.tile_size - self.frames.overlap + return (num_frames - 1 + frame_stride - 1) // frame_stride + + +# Either size-based (single-GPU VAE) or count-based (MGPU / explicit counts). +TilingConfig = TileSizeConfig | TileCountConfig + + +class AutoTiling: + """Sentinel: pipeline should recommend decode tiling (DiffVAE-aware / Conv default). + Distinct from ``None``, which means untiled decode. + """ + + __slots__ = () + + def __repr__(self) -> str: + return "AUTO_TILING" + + +AUTO_TILING = AutoTiling() + +# Pipeline ``tiling_config`` argument: explicit layout, auto-recommend, or untiled. +PipelineTiling = TilingConfig | AutoTiling | None + + +def _assert_video_on_vae_grid( + scale_factors: SpatioTemporalScaleFactors, + video_shape: VideoPixelShape, +) -> bool: + """Raise if ``video_shape`` is incompatible with the VAE ``scale_factors`` grid. + Returns whether the temporal axis is known (``video_shape.frames > 0``). When + False, callers skip frames-axis checks (duration not yet resolved). + """ + if scale_factors.time < 1 or scale_factors.height < 1 or scale_factors.width < 1: + raise ValueError(f"scale_factors must be >= 1 on each axis, got {scale_factors}") + if video_shape.height < 1 or video_shape.width < 1: + raise ValueError(f"video_shape height/width must be >= 1, got {video_shape.height}x{video_shape.width}") + if video_shape.height % scale_factors.height != 0: + raise ValueError(f"video height {video_shape.height} must be divisible by scale {scale_factors.height}") + if video_shape.width % scale_factors.width != 0: + raise ValueError(f"video width {video_shape.width} must be divisible by scale {scale_factors.width}") + if video_shape.frames <= 0: + return False + if (video_shape.frames - 1) % scale_factors.time != 0: + raise ValueError(f"video frames {video_shape.frames} must satisfy (frames - 1) % {scale_factors.time} == 0") + return True + + +def _validate_size_axis(cfg: DimensionSizeConfig, factor: int, axis_name: str) -> None: + """Pixel/frame size-axis legality vs VAE ``factor``.""" + if not cfg.is_tiled(): + return + min_size = 2 * factor + if cfg.tile_size < min_size: + raise ValueError(f"{axis_name}.tile_size must be at least {min_size}, got {cfg.tile_size}") + if cfg.tile_size % factor != 0: + raise ValueError(f"{axis_name}.tile_size must be divisible by {factor}, got {cfg.tile_size}") + if cfg.overlap % factor != 0: + raise ValueError(f"{axis_name}.overlap must be divisible by {factor}, got {cfg.overlap}") + + +def _validate_count_axis(cfg: DimensionTilingConfig, latent_extent: int, axis_name: str) -> None: + """Latent count-axis legality vs latent ``extent``.""" + if not cfg.is_tiled(): + return + if cfg.num_tiles > latent_extent: + raise ValueError(f"{axis_name}.num_tiles {cfg.num_tiles} exceeds latent {axis_name} extent {latent_extent}") + # split_by_count requires overlap < tile_size; tile_size grows with extent/n. + max_overlap = latent_extent - cfg.num_tiles + if cfg.overlap > max_overlap: + raise ValueError( + f"{axis_name}.overlap {cfg.overlap} exceeds latent bound {max_overlap} " + f"for extent {latent_extent} with {cfg.num_tiles} tiles" + ) + + +def _validate_overlap( + tiling_config: TilingConfig, + *, + min_overlap_frames: int, + min_overlap_pixels: int, +) -> None: + """Raise if any tiled ``TileSizeConfig`` axis overlap is below the given floors.""" + if not isinstance(tiling_config, TileSizeConfig): + return + + for axis_name, cfg, recommended, unit in ( + ("frames", tiling_config.frames, min_overlap_frames, "frames"), + ("height", tiling_config.height, min_overlap_pixels, "px"), + ("width", tiling_config.width, min_overlap_pixels, "px"), + ): + if cfg.is_tiled() and cfg.overlap < recommended: + raise ValueError(f"{axis_name} overlap {cfg.overlap} {unit} is below the required {recommended} {unit}.") + + +def balanced_tile_split(num_tiles: int) -> tuple[int, int]: + """Factor ``num_tiles`` into ``(small, large)`` as square as possible. + ``small`` is the largest divisor not exceeding the square root, so + ``small * large == num_tiles`` and ``small <= large``. E.g. 2 -> (1, 2), + 4 -> (2, 2), 8 -> (2, 4), 16 -> (4, 4). The caller decides which tiled + dimension gets which factor. + """ + if num_tiles < 1: + raise ValueError(f"num_tiles must be >= 1, got {num_tiles}") + small = next(d for d in range(math.isqrt(num_tiles), 0, -1) if num_tiles % d == 0) + return small, num_tiles // small diff --git a/telefuser/models/ltx25/diff_vae/transformer/__init__.py b/telefuser/models/ltx25/diff_vae/transformer/__init__.py new file mode 100644 index 0000000..3b8db2d --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/__init__.py @@ -0,0 +1,18 @@ +"""Eager neighborhood-attention building blocks for the isolated DiffVAE.""" + +from .attention import NeighborhoodAttention3D +from .blocks import DiffusionNABlock, NABlock +from .combined.block import CombinedDiffusionNABlock +from .config import DiffVAEMode +from .layers import AdaLNZero, ChannelLinear, LinearPixelShuffleUpsample + +__all__ = [ + "AdaLNZero", + "ChannelLinear", + "CombinedDiffusionNABlock", + "DiffVAEMode", + "DiffusionNABlock", + "LinearPixelShuffleUpsample", + "NABlock", + "NeighborhoodAttention3D", +] diff --git a/telefuser/models/ltx25/diff_vae/transformer/attention.py b/telefuser/models/ltx25/diff_vae/transformer/attention.py new file mode 100644 index 0000000..7a70a69 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/attention.py @@ -0,0 +1,150 @@ +"""3D Neighborhood Attention via NATTEN + absolute RoPE prelude. +Parameter shell shared by det ``NABlock`` and both diff-attn roles. +Diffusion AdaLN residuals live in pathway packages (each owns its RoPE); +det stages use ``det_attn_rope`` from :meth:`NeighborhoodAttention3D.forward`. +``attention_function`` selects the NA backend (NATTEN, Triton/eager fallback, or CuTe DSL). +""" + +from __future__ import annotations + +from typing import Protocol + +import torch +from torch import nn + +from telefuser.models.ltx25.diff_vae.transformer.det_attn_rope import det_qkv_rope +from telefuser.models.ltx25.diff_vae.transformer.qkv import QKVProjections +from telefuser.models.ltx25.diff_vae.transformer.rope_math import ( + DEFAULT_ABS_ROPE_NUM_TILES, + default_rope_dim_split, + rope_inv_freqs, +) +from telefuser.ops.neighborhood_attention import natten_available, neighborhood_attention_3d + + +class NAAttentionCallable(Protocol): + """A windowed 3D neighborhood-attention backend. + Q/K/V arrive as ``(B, T, H, W, NH, HD)``, already normed, scaled and RoPE'd; + the return is ``(B, T, H, W, NH*HD)`` or anything reshapeable to it. The owning + module is passed so a backend can read configuration (``kernel_size``, softmax + bound, …). Backend-specific settings (NATTEN's kernel pin) live on the callable. + """ + + def __call__( + self, + attn: NeighborhoodAttention3D, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> torch.Tensor: ... + + +class NattenAttention(NAAttentionCallable): + """``natten.na3d``, the default backend. + ``backend`` pins ``na3d``'s own kernel choice (e.g. ``"cutlass-fna"``); ``None`` + leaves NATTEN's auto-pick (hopper-fna on H100, etc.). + """ + + def __init__(self, backend: str | None = None) -> None: + self._backend = backend + + def __call__( + self, + attn: NeighborhoodAttention3D, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> torch.Tensor: + return neighborhood_attention_3d(q, k, v, kernel_size=attn.kernel_size, backend=self._backend) + + +class NeighborhoodAttention3D(nn.Module): + """3D Neighborhood Attention with absolute RoPE + pluggable NA backend. + Q/K receive absolute RoPE; attention is ``attention_function`` (NATTEN by + default; Triton or eager SDPA when natten is missing; CuTe DSL via + DiffVAE BLACKWELL_DSL install). Relative gather-based NA is not used as a + production gather path on this branch. + NATTEN shifts its window inward at grid boundaries instead of + clamp-and-mask; interior positions match the gather reference closely, + boundary positions may differ slightly. + """ + + def __init__( + self, + dim: int, + kernel_size: tuple[int, int, int], + head_dim: int = 64, + rope_dim_split: tuple[int, int, int] | None = None, + rope_base: float = 10000.0, + ) -> None: + super().__init__() + assert dim % head_dim == 0, f"dim={dim} not divisible by head_dim={head_dim}" + self.dim = dim + self.num_heads = dim // head_dim + self.head_dim = head_dim + self.kernel_size = tuple(kernel_size) + self.scale = head_dim**-0.5 + + if rope_dim_split is None: + rope_dim_split = default_rope_dim_split(head_dim) + assert sum(rope_dim_split) == head_dim, f"rope_dim_split={rope_dim_split} must sum to head_dim={head_dim}" + self.rope_dim_split = rope_dim_split + self.rope_base = rope_base + self.rope_num_tiles = DEFAULT_ABS_ROPE_NUM_TILES + self.rope_compute_dtype = torch.float32 + # Kept for the chunked opaque residual (string arg); callable is the swap surface. + self.natten_backend: str | None = None + self.attention_function: NAAttentionCallable = NattenAttention() + + self.register_buffer("rope_inv_t", rope_inv_freqs(rope_dim_split[0], rope_base), persistent=False) + self.register_buffer("rope_inv_h", rope_inv_freqs(rope_dim_split[1], rope_base), persistent=False) + self.register_buffer("rope_inv_w", rope_inv_freqs(rope_dim_split[2], rope_base), persistent=False) + + self.qkv = QKVProjections(dim) + self.proj = nn.Linear(dim, dim, bias=True) + self.q_norm = nn.RMSNorm(head_dim, eps=1e-6) + self.k_norm = nn.RMSNorm(head_dim, eps=1e-6) + + # W-chunking configuration (consumed by ``chunked.attn``). + self.w_chunks = 1 # 1 = no chunking + + def project_qkv(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Q/K/V as owned contiguous ``(B,T,H,W,NH,HD)`` tensors.""" + batch, t, h, w, _ = x.shape + q, k, v = self.qkv(x) + shape = (batch, t, h, w, self.num_heads, self.head_dim) + return q.view(shape), k.view(shape), v.view(shape) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Det-stage NA: opaque abs-RoPE via ``det_attn_rope`` + ``attention_function``. + ``x``/output: (B, T, H, W, C) — channels-last. RoPE positions are local + 0-based (see ``det_attn_rope`` module docstring for why that is + equivalent under tiled decode). + """ + batch, t, h, w, _ = x.shape + kt, kh, kw = self.kernel_size + if t < kt or h < kh or w < kw: + raise ValueError( + f"3D neighborhood attention requires spatial dims >= kernel_size; " + f"got (T,H,W)=({t},{h},{w}) vs kernel={self.kernel_size}" + ) + + q, k, v = det_qkv_rope(self, x) + # natten's CUTLASS kernel silently produces wrong output if inputs are non-contiguous. + q, k, v = q.contiguous(), k.contiguous(), v.contiguous() + out = self.attention_function(self, q, k, v) + out = out.reshape(batch, t, h, w, self.dim) + return self.proj(out) + + +def configure_w_chunks(module_root: nn.Module, w_chunks: int = 1) -> None: + """Set W-chunking on ``NeighborhoodAttention3D`` under ``module_root``. + When ``w_chunks > 1``, also sets ``rope_num_tiles=1`` so ``chunked.attn`` + owns the W axis (RoPE W-tiling would double-split). Pass only the diffusion + residual subtree — det-stage attention must keep its default RoPE tiling. + """ + for module in module_root.modules(): + if isinstance(module, NeighborhoodAttention3D): + module.w_chunks = w_chunks + if w_chunks > 1: + module.rope_num_tiles = 1 diff --git a/telefuser/models/ltx25/diff_vae/transformer/blocks.py b/telefuser/models/ltx25/diff_vae/transformer/blocks.py new file mode 100644 index 0000000..f3c31ea --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/blocks.py @@ -0,0 +1,83 @@ +"""NABlock and DiffusionNABlock parameter shells for DiffVAE. +Pathway subclasses live in ``chunked/`` and ``combined/``; ``apply`` installs +them via ``__class__`` swap (same pattern as ``Fp8CastLinear``). The shell owns +weights + shared AdaLN helpers only — no pathway forward. +""" + +from __future__ import annotations + +import torch +from torch import nn + +from telefuser.models.ltx25.diff_vae.transformer.attention import NeighborhoodAttention3D +from telefuser.models.ltx25.diff_vae.transformer.layers import AdaLNZero +from telefuser.models.ltx25.diff_vae.transformer.swiglu import SwiGLU, plain_mlp + +__all__ = [ + "DiffusionNABlock", + "NABlock", +] + + +class NABlock(nn.Module): + """Pre-norm transformer block: NA -> SwiGLU MLP with residual adds.""" + + def __init__( + self, + dim: int, + kernel_size: tuple[int, int, int], + head_dim: int = 64, + mlp_ratio: float = 4.0, + rope_dim_split: tuple[int, int, int] | None = None, + ) -> None: + super().__init__() + self.norm1 = nn.RMSNorm(dim, eps=1e-6) + self.attn = NeighborhoodAttention3D(dim, kernel_size, head_dim=head_dim, rope_dim_split=rope_dim_split) + self.norm2 = nn.RMSNorm(dim, eps=1e-6) + hidden = (int(dim * mlp_ratio) + 15) // 16 * 16 + self.mlp = SwiGLU(dim, hidden) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Channels-last in/out: (B, T, H, W, C).""" + x = x + self.attn(self.norm1(x)) + x = plain_mlp(x, self.mlp, self.norm2, self.mlp.tile) + return x + + +class DiffusionNABlock(nn.Module): + """Parameter shell for diffusion NA + SwiGLU with shared AdaLN-Zero. + Mode-specific subclasses (:class:`~ltx_core.model.video_vae.transformer.combined.block.CombinedDiffusionNABlock`, + :class:`~ltx_core.model.video_vae.transformer.chunked.block.ChunkedDiffusionNABlock`) + are installed via ModuleOps ``__class__`` swap and own the forward path. + Not a ``Protocol``: must be a concrete ``nn.Module`` so checkpoint load and + ``__class__`` swap keep one parameter identity. + """ + + def __init__( + self, + dim: int, + kernel_size: tuple[int, int, int], + context_channels: int, + head_dim: int = 64, + mlp_ratio: float = 4.0, + rope_dim_split: tuple[int, int, int] | None = None, + ) -> None: + super().__init__() + self.context_channels = context_channels + self.context_proj = nn.Linear(context_channels, dim, bias=True) + self.scale_shift_table = nn.Parameter(torch.zeros(AdaLNZero.NUM_CHUNKS, dim)) + + self.norm1 = nn.RMSNorm(dim, eps=1e-6) + self.attn = NeighborhoodAttention3D(dim, kernel_size, head_dim=head_dim, rope_dim_split=rope_dim_split) + self.norm2 = nn.RMSNorm(dim, eps=1e-6) + hidden = (int(dim * mlp_ratio) + 15) // 16 * 16 + self.mlp = SwiGLU(dim, hidden) + self.attn.proj.reset_parameters() + + def _modulation( + self, modulation: tuple[torch.Tensor, ...] + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + scale_msa, shift_msa, _, scale_mlp, shift_mlp, _, _ = [ + modulation[i] + self.scale_shift_table[i].view(1, 1, 1, 1, -1) for i in range(AdaLNZero.NUM_CHUNKS) + ] + return scale_msa, shift_msa, scale_mlp, shift_mlp diff --git a/telefuser/models/ltx25/diff_vae/transformer/chunked/__init__.py b/telefuser/models/ltx25/diff_vae/transformer/chunked/__init__.py new file mode 100644 index 0000000..42960f6 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/chunked/__init__.py @@ -0,0 +1 @@ +"""Chunked DiffVAE pathway: deferred inject, W-chunked residual attn, modulating MLP.""" diff --git a/telefuser/models/ltx25/diff_vae/transformer/chunked/attn.py b/telefuser/models/ltx25/diff_vae/transformer/chunked/attn.py new file mode 100644 index 0000000..86b44f0 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/chunked/attn.py @@ -0,0 +1,420 @@ +"""Chunked* diffusion AdaLN residual attention (W-loop + halo + in-slab RoPE). +Owns the fixed-extent W-chunk residual and its in-slab abs-RoPE (``w_pos``). +Does **not** call ``combined.attn``; shares only NA module weights (+ ``rope_math``). +NA goes through ``attn.attention_function`` (same pluggable backend as combined / det). +The W-chunk loop stays behind opaque ``ltx_core::na_residual_w_chunked`` so Dynamo +does not unroll it; the callable is bound via a contextvar around the op call +(custom_op schemas cannot take Python callables as args). +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from contextvars import ContextVar +from types import SimpleNamespace + +import torch +import torch.nn.functional as F +from torch import nn + +from telefuser.models.ltx25.diff_vae.transformer.attention import ( + NAAttentionCallable, + NattenAttention, + NeighborhoodAttention3D, + natten_available, +) +from telefuser.models.ltx25.diff_vae.transformer.rope_math import ( + h_positions, + rot_abs_axis_impl, + t_positions, +) + +# Bound by ``chunked()`` for the duration of the opaque custom_op. Dynamo only sees +# the op call site; the Python body reads these at runtime. +_CHUNKED_ATTENTION_FUNCTION: ContextVar[NAAttentionCallable | None] = ContextVar( + "ltx_chunked_attention_function", + default=None, +) +_CHUNKED_ATTN_MODULE: ContextVar[NeighborhoodAttention3D | None] = ContextVar( + "ltx_chunked_attn_module", + default=None, +) + + +@contextmanager +def _bound_chunked_attention( + attention_function: NAAttentionCallable, + attn: NeighborhoodAttention3D, +) -> Iterator[None]: + t_fn = _CHUNKED_ATTENTION_FUNCTION.set(attention_function) + t_mod = _CHUNKED_ATTN_MODULE.set(attn) + try: + yield + finally: + _CHUNKED_ATTENTION_FUNCTION.reset(t_fn) + _CHUNKED_ATTN_MODULE.reset(t_mod) + + +def _resolve_chunked_attention() -> tuple[NAAttentionCallable, NeighborhoodAttention3D]: + fn = _CHUNKED_ATTENTION_FUNCTION.get() + mod = _CHUNKED_ATTN_MODULE.get() + if fn is None or mod is None: + raise RuntimeError( + "ltx_core::na_residual_w_chunked requires attention_function binding; " + "call via chunked(), not the custom_op directly" + ) + return fn, mod + + +def _apply_in_slab_abs_rope( + x: torch.Tensor, + rope_split: tuple[int, int, int], + inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + w_pos: torch.Tensor, + *, + compute_dtype: torch.dtype, +) -> torch.Tensor: + """Single-extent abs-RoPE via raw rot (no RoPE tile loop, no nested region). + Used inside the W-chunk residual body: the chunk is already one W slab. + Nested-in-nested breaks FakeTensor mode matching under Inductor. + ``w_pos`` must be length ``x.shape[3]`` (global W coords for this slab). + """ + d_t, d_h, _ = rope_split + inv_t, inv_h, inv_w = inv_freqs + t = x.shape[1] + h = x.shape[2] + xt = rot_abs_axis_impl(x[..., :d_t], t_positions(t, x.device), inv_t, axis=1, compute_dtype=compute_dtype) + xh = rot_abs_axis_impl( + x[..., d_t : d_t + d_h], + h_positions(h, x.device), + inv_h, + axis=2, + compute_dtype=compute_dtype, + ) + xw = rot_abs_axis_impl(x[..., d_t + d_h :], w_pos, inv_w, axis=3, compute_dtype=compute_dtype) + return torch.cat([xt, xh, xw], dim=-1) + + +def _attn_on_chunk_impl( # noqa: PLR0913 + attention_function: NAAttentionCallable, + attn: NeighborhoodAttention3D | SimpleNamespace, + x_chunk: torch.Tensor, + w_pos: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + norm_weight: torch.Tensor, + q_w: torch.Tensor, + q_b: torch.Tensor, + k_w: torch.Tensor, + k_b: torch.Tensor, + v_w: torch.Tensor, + v_b: torch.Tensor, + q_norm_w: torch.Tensor, + k_norm_w: torch.Tensor, + proj_w: torch.Tensor, + proj_b: torch.Tensor, + inv_t: torch.Tensor, + inv_h: torch.Tensor, + inv_w: torch.Tensor, + d_t: int, + d_h: int, + d_w: int, + kt: int, + kh: int, + kw: int, + num_heads: int, + head_dim: int, + dim: int, + attn_scale: float, + rope_num_tiles: int, + rope_compute_bf16: bool, +) -> torch.Tensor: + """One W-slab: rms_norm + modulate + QKV / in-slab RoPE / ``attention_function`` / proj.""" + del rope_num_tiles, kt, kh, kw # chunk is already one W slab; kernel_size lives on ``attn`` + if not natten_available() and type(attention_function) is NattenAttention: + raise ImportError("natten is required for W-chunked residual attention") + + batch, t, h, ext_w, _ = x_chunk.shape + eps = 1e-6 + rope_dtype = torch.bfloat16 if rope_compute_bf16 else torch.float32 + inv_freqs = (inv_t, inv_h, inv_w) + rope_split = (d_t, d_h, d_w) + + y = F.rms_norm(x_chunk, (dim,), norm_weight, eps) + y = y * (1.0 + scale) + shift + + head_shape = (batch, t, h, ext_w, num_heads, head_dim) + q = F.linear(y, q_w, q_b).view(head_shape) + q = F.rms_norm(q, (head_dim,), q_norm_w, eps) * attn_scale + q = _apply_in_slab_abs_rope(q, rope_split, inv_freqs, w_pos, compute_dtype=rope_dtype) + k = F.linear(y, k_w, k_b).view(head_shape) + k = F.rms_norm(k, (head_dim,), k_norm_w, eps) + k = _apply_in_slab_abs_rope(k, rope_split, inv_freqs, w_pos, compute_dtype=rope_dtype) + v = F.linear(y, v_w, v_b).view(head_shape) + del y + + out = attention_function(attn, q.contiguous(), k.contiguous(), v.contiguous()) # type: ignore[arg-type] + del q, k, v + return F.linear(out.reshape(batch, t, h, ext_w, dim), proj_w, proj_b) + + +@torch.compiler.disable +def _pack_residual_weights(attn: NeighborhoodAttention3D, norm: nn.RMSNorm, x: torch.Tensor) -> tuple: + """Pack NA + norm weights / RoPE meta for the opaque chunked residual op.""" + d_t, d_h, d_w = attn.rope_dim_split + kt, kh, kw = attn.kernel_size + inv_t = attn.rope_inv_t.to(device=x.device) + inv_h = attn.rope_inv_h.to(device=x.device) + inv_w = attn.rope_inv_w.to(device=x.device) + return ( + norm.weight, + attn.qkv.to_q.weight, + attn.qkv.to_q.bias, + attn.qkv.to_k.weight, + attn.qkv.to_k.bias, + attn.qkv.to_v.weight, + attn.qkv.to_v.bias, + attn.q_norm.weight, + attn.k_norm.weight, + attn.proj.weight, + attn.proj.bias, + inv_t, + inv_h, + inv_w, + d_t, + d_h, + d_w, + kt, + kh, + kw, + attn.num_heads, + attn.head_dim, + attn.dim, + float(attn.scale), + attn.rope_num_tiles, + attn.rope_compute_dtype == torch.bfloat16, + ) + + +def _run_w_chunked_residual( + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + packed: tuple, + w_chunks: int, + halo: int, + attn_fn: Callable[..., torch.Tensor], +) -> torch.Tensor: + """In-place ``x += attn(...)`` with fixed-extent W slabs. + Neighbor-sourced halos are copied when available. Missing true-boundary + left/right halo slots (volume edges) are edge-replicated — same spirit as + trailing-frame replicate — so NA never sees zero-padded W boundaries. + """ + _, _, _, w, c = x.shape + chunk_w = (w + w_chunks - 1) // w_chunks + extent = chunk_w + 2 * halo + left_halo: torch.Tensor | None = None + + for i in range(w_chunks): + core_start = i * chunk_w + core_end = min(w, (i + 1) * chunk_w) + core_len = core_end - core_start + + buf = x.new_zeros(*x.shape[:3], extent, c) + if i > 0: + assert left_halo is not None + lh = left_halo.shape[3] + buf[:, :, :, halo - lh : halo, :] = left_halo + buf[:, :, :, halo : halo + core_len, :] = x[:, :, :, core_start:core_end, :] + right_filled = 0 + if i + 1 < w_chunks: + right_end = min(w, core_end + halo) + right = x[:, :, :, core_end:right_end, :] + right_filled = right.shape[3] + buf[:, :, :, halo + core_len : halo + core_len + right_filled, :] = right + + # True W-boundary: replicate-pad missing halo slots (do not overwrite neighbors). + if i == 0 and halo > 0 and core_len > 0: + edge_l = buf[:, :, :, halo : halo + 1, :] + buf[:, :, :, :halo, :] = edge_l.expand(*x.shape[:3], halo, c) + missing_right = extent - (halo + core_len + right_filled) + if missing_right > 0 and core_len > 0: + edge_r = buf[:, :, :, halo + core_len - 1 : halo + core_len, :] + lo_r = halo + core_len + right_filled + buf[:, :, :, lo_r:extent, :] = edge_r.expand(*x.shape[:3], missing_right, c) + + if i + 1 < w_chunks: + take = min(halo, core_len) + left_halo = x[:, :, :, core_end - take : core_end, :].clone() + + w_pos = torch.arange(extent, device=x.device, dtype=torch.float32) + (core_start - halo) + out = attn_fn(buf, w_pos, scale, shift, *packed) + x[:, :, :, core_start:core_end, :].add_(out[:, :, :, halo : halo + core_len, :]) + + return x + + +@torch.library.custom_op("ltx_core::na_residual_w_chunked", mutates_args=("x",)) +def _na_residual_w_chunked_op( # noqa: PLR0913 + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + norm_weight: torch.Tensor, + q_w: torch.Tensor, + q_b: torch.Tensor, + k_w: torch.Tensor, + k_b: torch.Tensor, + v_w: torch.Tensor, + v_b: torch.Tensor, + q_norm_w: torch.Tensor, + k_norm_w: torch.Tensor, + proj_w: torch.Tensor, + proj_b: torch.Tensor, + inv_t: torch.Tensor, + inv_h: torch.Tensor, + inv_w: torch.Tensor, + d_t: int, + d_h: int, + d_w: int, + kt: int, + kh: int, + kw: int, + num_heads: int, + head_dim: int, + dim: int, + attn_scale: float, + rope_num_tiles: int, + rope_compute_bf16: bool, + w_chunks: int, +) -> None: + """Opaque in-place W-chunked residual — Dynamo does not unroll the chunk loop. + Reads ``attention_function`` / NA module from the contextvars set by ``chunked()``. + """ + attention_function, attn = _resolve_chunked_attention() + halo = kw // 2 + packed = ( + norm_weight, + q_w, + q_b, + k_w, + k_b, + v_w, + v_b, + q_norm_w, + k_norm_w, + proj_w, + proj_b, + inv_t, + inv_h, + inv_w, + d_t, + d_h, + d_w, + kt, + kh, + kw, + num_heads, + head_dim, + dim, + attn_scale, + rope_num_tiles, + rope_compute_bf16, + ) + + def _slab( + x_chunk: torch.Tensor, + w_pos: torch.Tensor, + scale_: torch.Tensor, + shift_: torch.Tensor, + *rest: object, + ) -> torch.Tensor: + return _attn_on_chunk_impl(attention_function, attn, x_chunk, w_pos, scale_, shift_, *rest) # type: ignore[arg-type] + + _run_w_chunked_residual(x, scale, shift, packed, w_chunks, halo, _slab) + + +@_na_residual_w_chunked_op.register_fake +def _na_residual_w_chunked_fake( # noqa: PLR0913 + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + norm_weight: torch.Tensor, + q_w: torch.Tensor, + q_b: torch.Tensor, + k_w: torch.Tensor, + k_b: torch.Tensor, + v_w: torch.Tensor, + v_b: torch.Tensor, + q_norm_w: torch.Tensor, + k_norm_w: torch.Tensor, + proj_w: torch.Tensor, + proj_b: torch.Tensor, + inv_t: torch.Tensor, + inv_h: torch.Tensor, + inv_w: torch.Tensor, + d_t: int, + d_h: int, + d_w: int, + kt: int, + kh: int, + kw: int, + num_heads: int, + head_dim: int, + dim: int, + attn_scale: float, + rope_num_tiles: int, + rope_compute_bf16: bool, + w_chunks: int, +) -> None: + del ( + x, + scale, + shift, + norm_weight, + q_w, + q_b, + k_w, + k_b, + v_w, + v_b, + q_norm_w, + k_norm_w, + proj_w, + proj_b, + inv_t, + inv_h, + inv_w, + d_t, + d_h, + d_w, + kt, + kh, + kw, + num_heads, + head_dim, + dim, + attn_scale, + rope_num_tiles, + rope_compute_bf16, + w_chunks, + ) + + +def chunked( + x: torch.Tensor, + attn: NeighborhoodAttention3D, + norm: nn.RMSNorm, + scale: torch.Tensor, + shift: torch.Tensor, +) -> torch.Tensor: + """In-place W-chunked ``x += NA(modulate(norm(x)))`` with in-slab RoPE. + Opaque to Dynamo via ``ltx_core::na_residual_w_chunked``; NA uses + ``attn.attention_function`` (bound for the op body via contextvar). + """ + if attn.w_chunks < 1: + raise ValueError(f"w_chunks must be >= 1, got {attn.w_chunks}") + packed = _pack_residual_weights(attn, norm, x) + with _bound_chunked_attention(attn.attention_function, attn): + _na_residual_w_chunked_op(x, scale, shift, *packed, attn.w_chunks) + return x diff --git a/telefuser/models/ltx25/diff_vae/transformer/chunked/block.py b/telefuser/models/ltx25/diff_vae/transformer/chunked/block.py new file mode 100644 index 0000000..b9186e3 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/chunked/block.py @@ -0,0 +1,70 @@ +"""ChunkedDiffusionNABlock: deferred inject + W-chunked attn + modulating MLP.""" + +from __future__ import annotations + +import torch + +from telefuser.models.ltx25.diff_vae.transformer.blocks import DiffusionNABlock +from telefuser.models.ltx25.diff_vae.transformer.chunked.attn import chunked as residual_attn +from telefuser.models.ltx25.diff_vae.transformer.chunked.context import deferred as inject_context +from telefuser.models.ltx25.diff_vae.transformer.chunked.mlp import residual_modulating_mlp +from telefuser.models.ltx25.diff_vae.transformer.layers import LinearPixelShuffleUpsample + + +class ChunkedDiffusionNABlock(DiffusionNABlock): + """Deferred-context diffusion block. + ``forward`` / ``forward_x_ctx``: in-place inject (eager) then attn+mlp. + ``forward_attn_mlp``: opaque void-mutate residual + SwiGLU — this is what + ChunkedCompile ``torch.compile``s (do not compile inject). + ``stage4_upsample`` is wired by ``apply`` to ``decoder.upsamples[3]`` so + inject runs sequential upsample Linear then ``context_proj``. + """ + + stage4_upsample: LinearPixelShuffleUpsample | None + + def forward_attn_mlp( + self, + x: torch.Tensor, + scale_msa: torch.Tensor, + shift_msa: torch.Tensor, + scale_mlp: torch.Tensor, + shift_mlp: torch.Tensor, + ) -> torch.Tensor: + """In-place chunked NA residual + modulating SwiGLU (compile target).""" + x = residual_attn(x, self.attn, self.norm1, scale_msa, shift_msa) + x = residual_modulating_mlp(x, self.mlp, self.norm2, scale_mlp, shift_mlp, self.mlp.tile) + return x + + def forward_x_ctx( + self, + x: torch.Tensor, + stage4_feat: torch.Tensor, + modulation: tuple[torch.Tensor, ...], + *, + drop_leading_frame: bool = True, + ) -> torch.Tensor: + if self.stage4_upsample is None: + raise RuntimeError("stage4_upsample not configured; apply DiffVAEConfig with block=CHUNKED") + + scale_msa, shift_msa, scale_mlp, shift_mlp = self._modulation(modulation) + upsample = self.stage4_upsample + x = inject_context( + x, + stage4_feat, + upsample.proj, + self.context_proj, + tuple(upsample.stride), # type: ignore[arg-type] + w_chunks=self.attn.w_chunks, + drop_leading_frame=drop_leading_frame, + ) + return self.forward_attn_mlp(x, scale_msa, shift_msa, scale_mlp, shift_mlp) + + def forward( + self, + x: torch.Tensor, + stage4_feat: torch.Tensor, + modulation: tuple[torch.Tensor, ...], + *, + drop_leading_frame: bool = True, + ) -> torch.Tensor: + return self.forward_x_ctx(x, stage4_feat, modulation, drop_leading_frame=drop_leading_frame) diff --git a/telefuser/models/ltx25/diff_vae/transformer/chunked/context.py b/telefuser/models/ltx25/diff_vae/transformer/chunked/context.py new file mode 100644 index 0000000..67130cf --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/chunked/context.py @@ -0,0 +1,88 @@ +"""Deferred stage-4 context inject: sequential upsample then context_proj (W-chunked). +Eager in-place Python W-loop — stays outside ChunkedCompile's ``forward_attn_mlp`` +graph so Inductor never sees inject. Runs ``upsamples[3].proj`` (pixel-shuffle) +then ``context_proj``; no fused Linear. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from einops import rearrange +from torch import nn + + +def _upsample_then_ctx( + feat: torch.Tensor, + up_w: torch.Tensor, + up_b: torch.Tensor, + ctx_w: torch.Tensor, + ctx_b: torch.Tensor | None, + stride: tuple[int, int, int], + *, + drop_leading_frame: bool, +) -> torch.Tensor: + """``context_proj(pixel_shuffle(upsample_proj(feat)))``.""" + p1, p2, p3 = stride + up = F.linear(feat, up_w, up_b) + up = rearrange( + up, + "b t h w (c p1 p2 p3) -> b (t p1) (h p2) (w p3) c", + p1=p1, + p2=p2, + p3=p3, + ) + if p1 == 2 and drop_leading_frame: + up = up[:, 1:, :, :, :] + return F.linear(up, ctx_w, ctx_b) + + +def deferred( + x: torch.Tensor, + stage4_feat: torch.Tensor, + upsample_proj: nn.Linear, + context_proj: nn.Linear, + stride: tuple[int, int, int], + *, + w_chunks: int, + drop_leading_frame: bool, +) -> torch.Tensor: + """Chunk stage-4 feat along W, upsample+``context_proj`` each slab, ``add_`` into ``x``. + Returns ``x`` (mutated). + """ + assert upsample_proj.bias is not None + p3 = stride[2] + up_w, up_b = upsample_proj.weight, upsample_proj.bias + ctx_w, ctx_b = context_proj.weight, context_proj.bias + + if w_chunks <= 1: + x.add_( + _upsample_then_ctx( + stage4_feat, + up_w, + up_b, + ctx_w, + ctx_b, + stride, + drop_leading_frame=drop_leading_frame, + ) + ) + return x + + feat_chunks = torch.chunk(stage4_feat, w_chunks, dim=3) + w_hi = x.shape[3] + lo = 0 + for feat_chunk in feat_chunks: + hi = min(w_hi, lo + feat_chunk.shape[3] * p3) + ctx = _upsample_then_ctx( + feat_chunk, + up_w, + up_b, + ctx_w, + ctx_b, + stride, + drop_leading_frame=drop_leading_frame, + ) + x[:, :, :, lo:hi, :].add_(ctx[:, :, :, : hi - lo, :]) + lo = hi + return x diff --git a/telefuser/models/ltx25/diff_vae/transformer/chunked/mlp.py b/telefuser/models/ltx25/diff_vae/transformer/chunked/mlp.py new file mode 100644 index 0000000..6c4ef35 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/chunked/mlp.py @@ -0,0 +1,178 @@ +"""Chunked pathway MLP: inplace fused ``x += swiglu(modulate(rms_norm(x)))``.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch import nn + +from telefuser.models.ltx25.diff_vae.tiling import DimensionInterval +from telefuser.models.ltx25.diff_vae.transformer.swiglu import ( + SwiGLUTileSpec, + _fused_gate_up_swiglu_triton, + _fused_up_mul_torch, + _fused_up_mul_triton, + _tile_from_op_args, + dual_gate_up_triton_eligible, + token_intervals, + triton_swiglu_available, +) + + +def _channel_affine_bc(scale: torch.Tensor, shift: torch.Tensor, dim: int) -> tuple[torch.Tensor, torch.Tensor]: + s = scale.reshape(-1, dim) + sh = shift.reshape(-1, dim) + if s.shape[0] != 1 or sh.shape[0] != 1: + if s.shape[0] != sh.shape[0]: + raise ValueError(f"scale/shift batch mismatch: {tuple(s.shape)} vs {tuple(sh.shape)}") + if s.shape[0] != 1: + raise ValueError( + f"swiglu residual modulated expects broadcastable channel affine " + f"(got scale {tuple(scale.shape)}); B>1 not supported" + ) + return s, sh + + +def _chunked_residual_modulated_impl( + x: torch.Tensor, + norm_weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + w_gate: torch.Tensor, + w_up: torch.Tensor, + w_down: torch.Tensor, + intervals: list[DimensionInterval], + *, + use_triton: bool, + eps: float = 1e-6, +) -> None: + """In-place: ``x += swiglu(modulate(rms_norm(x)))`` with peak O(chunk·hidden).""" + if use_triton and not triton_swiglu_available(): + raise RuntimeError("use_triton=True but Triton/CUDA is unavailable") + + dim = x.shape[-1] + if norm_weight.shape != (dim,): + raise ValueError(f"norm_weight shape {tuple(norm_weight.shape)} != ({dim},)") + x_flat = x.reshape(-1, dim) + if not x_flat.is_contiguous(): + raise ValueError("x must be contiguous channels-last so reshape(-1, dim) is a writable view") + n_tok = x_flat.shape[0] + hidden = w_gate.shape[0] + if w_up.shape[0] != hidden or w_gate.shape[1] != dim or w_up.shape[1] != dim: + raise ValueError( + f"weight shapes incompatible with x[..., {dim}]: gate={tuple(w_gate.shape)} up={tuple(w_up.shape)}" + ) + if w_down.shape[1] != hidden or w_down.shape[0] != dim: + raise ValueError(f"w_down {tuple(w_down.shape)} incompatible with hidden={hidden} dim={dim}") + if n_tok == 0: + return + + s, sh = _channel_affine_bc(scale, shift, dim) + max_chunk = max((iv.end - iv.start for iv in intervals), default=0) + workspace = torch.empty((max_chunk, hidden), device=x.device, dtype=x.dtype) + y_buf = torch.empty((max_chunk, dim), device=x.device, dtype=x.dtype) + out_buf = torch.empty((max_chunk, dim), device=x.device, dtype=x.dtype) + use_dual = bool(use_triton) and dual_gate_up_triton_eligible(x, w_gate, w_up) + w_gate_c = w_gate.contiguous() if use_dual else w_gate + w_up_c = w_up.contiguous() if use_triton else w_up + + for iv in intervals: + start, end = iv.start, iv.end + if start >= end: + continue + n = end - start + xc = x_flat[start:end] + y = y_buf[:n] + y.copy_(F.rms_norm(xc, (dim,), norm_weight, eps)) + y.mul_(1.0 + s).add_(sh) + + ws = workspace[:n] + if use_dual: + _fused_gate_up_swiglu_triton(y, w_gate_c, w_up_c, ws) + else: + torch.mm(y, w_gate.t(), out=ws) + F.silu(ws, inplace=True) + if use_triton: + _fused_up_mul_triton(y, w_up_c, ws) + else: + _fused_up_mul_torch(y, w_up, ws) + torch.mm(ws, w_down.t(), out=out_buf[:n]) + xc.add_(out_buf[:n]) + + +@torch.library.custom_op("ltx_core::swiglu_tiled_residual_modulated", mutates_args=("x",)) +def _swiglu_tiled_residual_modulated_op( + x: torch.Tensor, + norm_weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + w_gate: torch.Tensor, + w_up: torch.Tensor, + w_down: torch.Tensor, + tile_size: int, + num_tiles: int, + use_triton: bool, +) -> None: + """Opaque op: chunked ``x += swiglu(modulate(rms_norm(x)))`` (compile-friendly).""" + tile = _tile_from_op_args(tile_size, num_tiles) + n_tok = x.reshape(-1, x.shape[-1]).shape[0] + _chunked_residual_modulated_impl( + x, + norm_weight, + scale, + shift, + w_gate, + w_up, + w_down, + token_intervals(n_tok, tile), + use_triton=use_triton, + ) + + +@_swiglu_tiled_residual_modulated_op.register_fake +def _swiglu_tiled_residual_modulated_fake( + x: torch.Tensor, + norm_weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + w_gate: torch.Tensor, + w_up: torch.Tensor, + w_down: torch.Tensor, + tile_size: int, + num_tiles: int, + use_triton: bool, +) -> None: + del x, norm_weight, scale, shift, w_gate, w_up, w_down, tile_size, num_tiles, use_triton + + +def residual_modulating_mlp( + x: torch.Tensor, + mlp: nn.Module, + norm: nn.RMSNorm, + scale: torch.Tensor, + shift: torch.Tensor, + tile: SwiGLUTileSpec, +) -> torch.Tensor: + """Inplace fused ``x += swiglu(modulate(rms_norm(x)))`` (Chunked path).""" + if x.numel() == 0: + return x + if not x.is_contiguous(): + x = x.contiguous() + + tile_size = int(tile.tile_size) if tile.tile_size is not None else 0 + num_tiles = int(tile.num_tiles) if tile.num_tiles is not None else 0 + use_triton = triton_swiglu_available() and x.is_cuda + + _swiglu_tiled_residual_modulated_op( + x, + norm.weight, + scale, + shift, + mlp.w_gate.weight, + mlp.w_up.weight, + mlp.w_down.weight, + tile_size, + num_tiles, + bool(use_triton), + ) + return x diff --git a/telefuser/models/ltx25/diff_vae/transformer/combined/__init__.py b/telefuser/models/ltx25/diff_vae/transformer/combined/__init__.py new file mode 100644 index 0000000..2b3f6c3 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/combined/__init__.py @@ -0,0 +1 @@ +"""Combined DiffVAE pathway: context_and_x inject, full-volume residual attn, residual MLP.""" diff --git a/telefuser/models/ltx25/diff_vae/transformer/combined/attn.py b/telefuser/models/ltx25/diff_vae/transformer/combined/attn.py new file mode 100644 index 0000000..2bf7a8d --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/combined/attn.py @@ -0,0 +1,125 @@ +"""Combined* diffusion AdaLN residual attention (full-volume NA + nested RoPE). +Owns nested full-volume abs-RoPE for the Combined / ``w_chunks==1`` path. +Does not share a residual body with ``chunked`` — only the NA module weights. +""" + +from __future__ import annotations + +import torch +from torch import nn + +from telefuser.models.ltx25.diff_vae.transformer.attention import NeighborhoodAttention3D +from telefuser.models.ltx25.diff_vae.transformer.rope_math import ( + h_positions, + rot_abs_axis_impl, + t_positions, +) + +# ``nested_compile_region`` was introduced after some Torch versions TeleFuser +# still supports. Eager execution has the same semantics without the wrapper. +_nested_compile_region = getattr(torch.compiler, "nested_compile_region", lambda function: function) +_rot_abs_axis = _nested_compile_region(rot_abs_axis_impl) + + +def _apply_nested_abs_rope_slab( + x: torch.Tensor, + rope_split: tuple[int, int, int], + inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + *, + w_pos: torch.Tensor, + compute_dtype: torch.dtype, +) -> torch.Tensor: + """Rotate one W-extent with nested per-axis abs-RoPE.""" + d_t, d_h, _ = rope_split + inv_t, inv_h, inv_w = inv_freqs + t = x.shape[1] + h = x.shape[2] + xt = _rot_abs_axis(x[..., :d_t], t_positions(t, x.device), inv_t, axis=1, compute_dtype=compute_dtype) + xh = _rot_abs_axis( + x[..., d_t : d_t + d_h], + h_positions(h, x.device), + inv_h, + axis=2, + compute_dtype=compute_dtype, + ) + xw = _rot_abs_axis(x[..., d_t + d_h :], w_pos, inv_w, axis=3, compute_dtype=compute_dtype) + return torch.cat([xt, xh, xw], dim=-1) + + +def _apply_nested_full_volume_rope( + x: torch.Tensor, + rope_split: tuple[int, int, int], + inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + *, + num_tiles: int, + compute_dtype: torch.dtype, +) -> torch.Tensor: + """Fixed-``num_tiles`` W split + nested per-slab rotation (Dynamo-safe).""" + slabs = torch.chunk(x, num_tiles, dim=3) + w_off = 0 + parts: list[torch.Tensor] = [] + for slab in slabs: + w_slab = slab.shape[3] + w_pos = torch.arange(w_slab, dtype=torch.float32, device=x.device) + w_off + parts.append( + _apply_nested_abs_rope_slab( + slab, + rope_split, + inv_freqs, + w_pos=w_pos, + compute_dtype=compute_dtype, + ) + ) + w_off = w_off + w_slab + return torch.cat(parts, dim=3) + + +def _qkv_nested_rope(attn: NeighborhoodAttention3D, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Q/K/V proj + norm/scale + nested full-volume abs-RoPE.""" + q, k, v = attn.project_qkv(x) + q = attn.q_norm(q) * attn.scale + k = attn.k_norm(k) + inv_freqs = ( + attn.rope_inv_t.to(device=x.device), + attn.rope_inv_h.to(device=x.device), + attn.rope_inv_w.to(device=x.device), + ) + q = _apply_nested_full_volume_rope( + q, + attn.rope_dim_split, + inv_freqs, + num_tiles=attn.rope_num_tiles, + compute_dtype=attn.rope_compute_dtype, + ) + k = _apply_nested_full_volume_rope( + k, + attn.rope_dim_split, + inv_freqs, + num_tiles=attn.rope_num_tiles, + compute_dtype=attn.rope_compute_dtype, + ) + return q, k, v + + +def full( + x: torch.Tensor, + attn: NeighborhoodAttention3D, + norm: nn.RMSNorm, + scale: torch.Tensor, + shift: torch.Tensor, +) -> torch.Tensor: + """``x + NA(modulate(norm(x)))`` with nested full-volume abs-RoPE.""" + y = norm(x) * (1.0 + scale) + shift + batch, t, h, w, _ = y.shape + kt, kh, kw = attn.kernel_size + if t < kt or h < kh or w < kw: + raise ValueError( + f"3D neighborhood attention requires spatial dims >= kernel_size; " + f"got (T,H,W)=({t},{h},{w}) vs kernel={attn.kernel_size}" + ) + + q, k, v = _qkv_nested_rope(attn, y) + q, k, v = q.contiguous(), k.contiguous(), v.contiguous() + out = attn.attention_function(attn, q, k, v) + out = out.reshape(batch, t, h, w, attn.dim) + return x + attn.proj(out) diff --git a/telefuser/models/ltx25/diff_vae/transformer/combined/block.py b/telefuser/models/ltx25/diff_vae/transformer/combined/block.py new file mode 100644 index 0000000..eae42b7 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/combined/block.py @@ -0,0 +1,32 @@ +"""CombinedDiffusionNABlock: context_and_x inject + full-volume attn + residual MLP.""" + +from __future__ import annotations + +import torch + +from telefuser.models.ltx25.diff_vae.transformer.blocks import DiffusionNABlock +from telefuser.models.ltx25.diff_vae.transformer.combined.attn import full as residual_attn +from telefuser.models.ltx25.diff_vae.transformer.combined.context import combined as inject_context +from telefuser.models.ltx25.diff_vae.transformer.combined.mlp import residual_mlp + + +class CombinedDiffusionNABlock(DiffusionNABlock): + """Combined-context diffusion block: ``forward`` / ``forward_combined``.""" + + def forward_combined( + self, + context_and_x: torch.Tensor, + modulation: tuple[torch.Tensor, ...], + ) -> torch.Tensor: + scale_msa, shift_msa, scale_mlp, shift_mlp = self._modulation(modulation) + x = inject_context(context_and_x, self.context_proj.weight, self.context_proj.bias) + x = residual_attn(x, self.attn, self.norm1, scale_msa, shift_msa) + x = residual_mlp(x, self.mlp, self.norm2, scale_mlp, shift_mlp, self.mlp.tile) + return x + + def forward( + self, + context_and_x: torch.Tensor, + modulation: tuple[torch.Tensor, ...], + ) -> torch.Tensor: + return self.forward_combined(context_and_x, modulation) diff --git a/telefuser/models/ltx25/diff_vae/transformer/combined/context.py b/telefuser/models/ltx25/diff_vae/transformer/combined/context.py new file mode 100644 index 0000000..7219574 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/combined/context.py @@ -0,0 +1,21 @@ +"""Combined context residual: project context half of ``context_and_x`` into ``x``.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + + +def combined( + context_and_x: torch.Tensor, + w_proj: torch.Tensor, + b_proj: torch.Tensor | None, +) -> torch.Tensor: + """Split ``context_and_x`` via ``w_proj.shape[1]`` (context channels), add projected ctx. + Returns the updated ``x`` half (not the full concatenated buffer). + ``w_proj`` is ``context_proj.weight`` with shape ``(dim, context_channels)``. + """ + context_channels = w_proj.shape[1] + latent_context = context_and_x[..., :context_channels] + x = context_and_x[..., context_channels:] + return x + F.linear(latent_context, w_proj, b_proj) diff --git a/telefuser/models/ltx25/diff_vae/transformer/combined/mlp.py b/telefuser/models/ltx25/diff_vae/transformer/combined/mlp.py new file mode 100644 index 0000000..b3be7b4 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/combined/mlp.py @@ -0,0 +1,24 @@ +"""Combined pathway MLP: out-of-place AdaLN SwiGLU residual.""" + +from __future__ import annotations + +import torch +from torch import nn + +from telefuser.models.ltx25.diff_vae.transformer.layers import modulate +from telefuser.models.ltx25.diff_vae.transformer.swiglu import SwiGLUTileSpec, swiglu_tiled + + +def residual_mlp( + x: torch.Tensor, + mlp: nn.Module, + norm: nn.RMSNorm, + scale: torch.Tensor, + shift: torch.Tensor, + tile: SwiGLUTileSpec, +) -> torch.Tensor: + """Combined*: ``x + swiglu_tiled(modulate(norm(x), scale, shift))``.""" + y = modulate(norm(x), scale, shift) + if y.numel() == 0: + return x + return x + swiglu_tiled(y, mlp.w_gate.weight, mlp.w_up.weight, mlp.w_down.weight, tile) diff --git a/telefuser/models/ltx25/diff_vae/transformer/compiling.py b/telefuser/models/ltx25/diff_vae/transformer/compiling.py new file mode 100644 index 0000000..306425b --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/compiling.py @@ -0,0 +1,39 @@ +"""``torch.compile`` setup for the isolated LTX-2.5 DiffVAE decoder.""" + +from __future__ import annotations + +from typing import Any, Callable + +import torch + +from telefuser.models.ltx25.diff_vae.transformer.chunked.block import ChunkedDiffusionNABlock +from telefuser.ops.neighborhood_attention import configure_neighborhood_attention_kv_parallelism + + +def compile_diffusion_decoder(decoder: torch.nn.Module) -> torch.nn.Module: + """Compile the upstream CHUNKED_COMPILE stage-5 residual methods. + + Context injection remains eager and in-place. Compiling only + ``forward_attn_mlp`` keeps the peak allocation near the eager chunked path. + """ + if not hasattr(torch, "compile"): + raise RuntimeError("CHUNKED_COMPILE requires PyTorch 2.0 or newer") + + configure_neighborhood_attention_kv_parallelism(False) + compile_kwargs: dict[str, Any] = { + "mode": None, + "backend": "inductor", + "fullgraph": False, + "dynamic": None, + } + + def _compile(function: Callable[..., Any]) -> Callable[..., Any]: + with torch._dynamo.config.patch(inline_inbuilt_nn_modules=True, cache_size_limit=256): # type: ignore[attr-defined] + return torch.compile(function, **compile_kwargs) + + for block in decoder.diff_blocks: # type: ignore[attr-defined] + if isinstance(block, ChunkedDiffusionNABlock): + block.forward_attn_mlp = _compile(block.forward_attn_mlp) # type: ignore[method-assign] + + decoder.mark_dynamic_shapes = True # type: ignore[attr-defined] + return decoder diff --git a/telefuser/models/ltx25/diff_vae/transformer/config.py b/telefuser/models/ltx25/diff_vae/transformer/config.py new file mode 100644 index 0000000..2489fc3 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/config.py @@ -0,0 +1,110 @@ +"""DiffVAE decode presets and resolved install recipes (pure data — no module I/O).""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum + +# Default eager DiffVAE pin: lower VRAM than hopper-fna TokPerm. +_CUTLASS_FNA_BACKEND = "cutlass-fna" + +# Fixed W split for Chunked* modes (matrix sweet spot at ~1088x1920). +_CHUNKED_W_CHUNKS = 4 + + +class NAttentionKind(Enum): + """Which neighborhood-attention backend ``apply`` installs on NA modules.""" + + NATTEN = "natten" + BLACKWELL_DSL = "blackwell_dsl" + TRITON = "triton" + EAGER_SDPA = "eager_sdpa" + + +class DiffVAEBlockKind(Enum): + """Which DiffusionNABlock subclass ``apply`` installs via ``__class__`` swap.""" + + COMBINED = "combined" + CHUNKED = "chunked" + BLACKWELL_DSL = "blackwell_dsl" + + +@dataclass(frozen=True, slots=True) +class DiffVAEConfig: + """Resolved install recipe — facts ``apply`` reads, not runtime forward knobs. + ``block`` selects the pathway (chunked vs combined vs Blackwell DSL); class swap + implies deferred inject + W-chunked attn + modulating MLP, combined context + + full-volume attn + residual MLP, or deferred stage-4 + fused CuTe DSL block. + """ + + block: DiffVAEBlockKind + w_chunks: int + natten_backend: str | None + attention: NAttentionKind + compile_blocks: bool + compile_det_stages: bool + + +class DiffVAEMode(Enum): + """User-facing DiffVAE decode presets. + Relative performance (order-of-magnitude; hardware varies — no absolute + timings or VRAM figures): + - **Compile:** ``chunked_compile`` roughly ~2x faster to compile than + ``combined_compile`` (det stages not compiled). + - **Warm runtime:** relative to ``combined_compile`` (fastest): + ``chunked_compile`` roughly ~1.4x slower, ``chunked_eager`` roughly + ~2-2.5x slower. + - **Peak VRAM:** ``chunked_*`` roughly ~½ of ``combined_compile``. + - **BLACKWELL_DSL:** datacenter Blackwell CuTe DSL NA + fused stage-5 + (deferred stage-4 inputs; upsample+context_proj inside the fused kernel). + CHUNKED_* always use deferred stage-4 + w_chunks=4. COMBINED_COMPILE is + combined context + w_chunks=1 (best runtime, highest memory) and **requires + natten**. Chunked modes fall back to Triton/eager when natten is missing. + """ + + CHUNKED_EAGER = "chunked_eager" + CHUNKED_COMPILE = "chunked_compile" + COMBINED_COMPILE = "combined_compile" + BLACKWELL_DSL = "blackwell_dsl" + + def resolve(self) -> DiffVAEConfig: + """Expand this preset into a concrete install recipe. + This is the only place preset → recipe expansion lives. + """ + if self is DiffVAEMode.CHUNKED_EAGER: + return DiffVAEConfig( + block=DiffVAEBlockKind.CHUNKED, + w_chunks=_CHUNKED_W_CHUNKS, + natten_backend=_CUTLASS_FNA_BACKEND, + attention=NAttentionKind.NATTEN, + compile_blocks=False, + compile_det_stages=False, + ) + if self is DiffVAEMode.CHUNKED_COMPILE: + return DiffVAEConfig( + block=DiffVAEBlockKind.CHUNKED, + w_chunks=_CHUNKED_W_CHUNKS, + natten_backend=None, + attention=NAttentionKind.NATTEN, + compile_blocks=True, + compile_det_stages=False, + ) + if self is DiffVAEMode.COMBINED_COMPILE: + return DiffVAEConfig( + block=DiffVAEBlockKind.COMBINED, + w_chunks=1, + natten_backend=None, + attention=NAttentionKind.NATTEN, + compile_blocks=True, + compile_det_stages=True, + ) + if self is DiffVAEMode.BLACKWELL_DSL: + return DiffVAEConfig( + block=DiffVAEBlockKind.BLACKWELL_DSL, + w_chunks=1, + natten_backend=None, + attention=NAttentionKind.BLACKWELL_DSL, + compile_blocks=True, + compile_det_stages=True, + ) + raise ValueError(f"Unknown DiffVAEMode: {self!r}") diff --git a/telefuser/models/ltx25/diff_vae/transformer/det_attn_rope.py b/telefuser/models/ltx25/diff_vae/transformer/det_attn_rope.py new file mode 100644 index 0000000..0b830e6 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/det_attn_rope.py @@ -0,0 +1,175 @@ +"""Opaque full-volume abs-RoPE for deterministic (pre-diffusion) NA. +Owns QKV + opaque ``custom_op`` packaging for det ``NA.forward``. Det stages +have differently shaped T/H/W; the opaque op keeps Dynamo from specializing +on each stage shape. Diffusion paths use ``diff_attn/`` RoPE — not this file. +No T/H/W origin/offset is threaded for tiled decode, and none is needed: +every attention call here is ``natten.na3d``, a local window with no +cross-tile tokens. Absolute-vs-local RoPE differs by a global phase that +cancels inside the attention softmax over that window, so the attention +output is unchanged. Since every tiled-decode call processes exactly one +tile in isolation, using each tile's local 0-based positions is identical +to using its true absolute origin. Absolute origin still matters for +whether a tile contains the latent's first frame (``drop_leading_frame``), +which is handled outside RoPE in the decoder stage / per-tile decode. +""" + +from __future__ import annotations + +import torch + +from telefuser.models.ltx25.diff_vae.transformer.rope_math import ( + DEFAULT_ABS_ROPE_NUM_TILES, + h_positions, + rot_abs_axis_impl, + t_positions, +) + + +def _apply_opaque_rope_slab( + x: torch.Tensor, + rope_split: tuple[int, int, int], + inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + *, + w_pos: torch.Tensor, + compute_dtype: torch.dtype, +) -> torch.Tensor: + """Rotate one W-extent with raw abs-RoPE (runs inside the opaque op).""" + d_t, d_h, _ = rope_split + inv_t, inv_h, inv_w = inv_freqs + t = x.shape[1] + h = x.shape[2] + xt = rot_abs_axis_impl(x[..., :d_t], t_positions(t, x.device), inv_t, axis=1, compute_dtype=compute_dtype) + xh = rot_abs_axis_impl( + x[..., d_t : d_t + d_h], + h_positions(h, x.device), + inv_h, + axis=2, + compute_dtype=compute_dtype, + ) + xw = rot_abs_axis_impl(x[..., d_t + d_h :], w_pos, inv_w, axis=3, compute_dtype=compute_dtype) + return torch.cat([xt, xh, xw], dim=-1) + + +def _apply_opaque_tiled_rope( + x: torch.Tensor, + rope_split: tuple[int, int, int], + inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + *, + num_tiles: int, + compute_dtype: torch.dtype, +) -> torch.Tensor: + """Fixed-``num_tiles`` W split + per-slab rotation (body of the opaque op).""" + slabs = torch.chunk(x, num_tiles, dim=3) + w_off = 0 + parts: list[torch.Tensor] = [] + for slab in slabs: + w_slab = slab.shape[3] + w_pos = torch.arange(w_slab, dtype=torch.float32, device=x.device) + w_off + parts.append( + _apply_opaque_rope_slab( + slab, + rope_split, + inv_freqs, + w_pos=w_pos, + compute_dtype=compute_dtype, + ) + ) + w_off = w_off + w_slab + return torch.cat(parts, dim=3) + + +@torch.library.custom_op("ltx_core::abs_rope", mutates_args=()) +def _abs_rope_op( + x: torch.Tensor, + inv_t: torch.Tensor, + inv_h: torch.Tensor, + inv_w: torch.Tensor, + d_t: int, + d_h: int, + d_w: int, + num_tiles: int, + compute_dtype_is_bf16: bool, +) -> torch.Tensor: + """Opaque out-of-place abs-RoPE: Dynamo sees one node.""" + compute_dtype = torch.bfloat16 if compute_dtype_is_bf16 else torch.float32 + return _apply_opaque_tiled_rope( + x, + (d_t, d_h, d_w), + (inv_t, inv_h, inv_w), + num_tiles=num_tiles, + compute_dtype=compute_dtype, + ) + + +@_abs_rope_op.register_fake +def _abs_rope_fake( + x: torch.Tensor, + inv_t: torch.Tensor, + inv_h: torch.Tensor, + inv_w: torch.Tensor, + d_t: int, + d_h: int, + d_w: int, + num_tiles: int, + compute_dtype_is_bf16: bool, +) -> torch.Tensor: + del inv_t, inv_h, inv_w, d_t, d_h, d_w, num_tiles, compute_dtype_is_bf16 + return torch.empty(x.shape, device=x.device, dtype=x.dtype) + + +def _apply_opaque_abs_rope( + x: torch.Tensor, + rope_split: tuple[int, int, int], + inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + *, + num_tiles: int, + compute_dtype: torch.dtype, +) -> torch.Tensor: + if num_tiles < 1: + raise ValueError(f"num_tiles must be >= 1, got {num_tiles}") + if compute_dtype not in (torch.float32, torch.bfloat16): + raise ValueError(f"compute_dtype must be float32 or bfloat16, got {compute_dtype}") + d_t, d_h, d_w = rope_split + inv_t, inv_h, inv_w = inv_freqs + return _abs_rope_op( + x, + inv_t, + inv_h, + inv_w, + d_t, + d_h, + d_w, + num_tiles, + compute_dtype == torch.bfloat16, + ) + + +def det_qkv_rope(attn: object, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Q/K/V proj + norm/scale + opaque full-volume abs-RoPE.""" + q, k, v = attn.project_qkv(x) # type: ignore[attr-defined] + q = attn.q_norm(q) # type: ignore[attr-defined] + k = attn.k_norm(k) # type: ignore[attr-defined] + q = q * attn.scale # type: ignore[attr-defined] + + inv_freqs = ( + attn.rope_inv_t.to(device=x.device), # type: ignore[attr-defined] + attn.rope_inv_h.to(device=x.device), # type: ignore[attr-defined] + attn.rope_inv_w.to(device=x.device), # type: ignore[attr-defined] + ) + num_tiles = getattr(attn, "rope_num_tiles", DEFAULT_ABS_ROPE_NUM_TILES) + compute_dtype = getattr(attn, "rope_compute_dtype", torch.float32) + q = _apply_opaque_abs_rope( + q, + attn.rope_dim_split, # type: ignore[attr-defined] + inv_freqs, + num_tiles=num_tiles, + compute_dtype=compute_dtype, + ) + k = _apply_opaque_abs_rope( + k, + attn.rope_dim_split, # type: ignore[attr-defined] + inv_freqs, + num_tiles=num_tiles, + compute_dtype=compute_dtype, + ) + return q, k, v diff --git a/telefuser/models/ltx25/diff_vae/transformer/fallback_na/__init__.py b/telefuser/models/ltx25/diff_vae/transformer/fallback_na/__init__.py new file mode 100644 index 0000000..b9d7141 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/fallback_na/__init__.py @@ -0,0 +1,154 @@ +"""DiffVAE NA fallbacks when natten is unavailable: Triton then eager SDPA. +Vendored from comfy-kitchen (Apache-2.0). Selection order for NATTEN-kind recipes:: + natten → (loud warning) Triton → eager tiled SDPA +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import torch + +from telefuser.models.ltx25.diff_vae.transformer.fallback_na.eager import na3d as eager_na3d + +if TYPE_CHECKING: + from telefuser.models.ltx25.diff_vae.transformer.attention import NAAttentionCallable, NeighborhoodAttention3D + +logger = logging.getLogger(__name__) + +_NO_NATTEN_WARNING = ( + "================================================================================\n" + "DiffVAE: natten is NOT installed. Falling back to a slower neighborhood-attention\n" + "backend (Triton if available, else pure-PyTorch tiled SDPA). This path is for\n" + "compatibility only — install natten for production DiffVAE decode:\n" + " uv sync --package ltx-core --extra natten\n" + "================================================================================" +) + + +def triton_na_available() -> bool: + """True when CUDA is available and the ``triton`` package imports cleanly. + Triton on Windows is supported via the community ``triton-windows`` builds + (https://github.com/triton-lang/triton-windows); there is no platform ban here. + Import failures beyond ``ImportError`` (e.g. ``OSError`` when libcuda is missing) + are treated as unavailable so hosts fall back to eager SDPA. + """ + if not torch.cuda.is_available(): + return False + try: + import triton # noqa: F401, PLC0415 + except (ImportError, OSError): + return False + return True + + +@torch.library.custom_op("ltx_core::na_attention_eager", mutates_args=()) +def na_attention_eager( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kernel_size: list[int], +) -> torch.Tensor: + """Opaque eager tiled-SDPA NA so Dynamo does not graph-break into the Python loop.""" + return eager_na3d(q, k, v, kernel_size=kernel_size, is_causal=None, scale=1.0) + + +@na_attention_eager.register_fake +def _na_attention_eager_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kernel_size: list[int], +) -> torch.Tensor: + del k, v, kernel_size + return torch.empty_like(q) + + +@torch.library.custom_op("ltx_core::na_attention_triton", mutates_args=()) +def na_attention_triton( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kernel_size: list[int], +) -> torch.Tensor: + """Opaque Triton NA launch so the surrounding block graph stays whole under compile.""" + from telefuser.models.ltx25.diff_vae.transformer.fallback_na.triton_na import ( # noqa: PLC0415 + na3d as triton_na3d, + ) + + return triton_na3d(q, k, v, kernel_size=kernel_size, is_causal=None, scale=1.0) + + +@na_attention_triton.register_fake +def _na_attention_triton_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kernel_size: list[int], +) -> torch.Tensor: + del k, v, kernel_size + return torch.empty_like(q) + + +class EagerSdpaAttention: + """Limited-workspace tiled SDPA NA (always available).""" + + def __call__( + self, + attn: NeighborhoodAttention3D, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> torch.Tensor: + if q.dtype != v.dtype or k.dtype != v.dtype: + q = q.to(dtype=v.dtype) + k = k.to(dtype=v.dtype) + return na_attention_eager(q, k, v, list(attn.kernel_size)) + + +class TritonNaAttention: + """Triton flash-style NA (hosts with CUDA + a working Triton install).""" + + def __call__( + self, + attn: NeighborhoodAttention3D, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> torch.Tensor: + if not triton_na_available(): + raise ImportError( + "Triton neighborhood attention requires CUDA and the triton package " + "(on Windows: https://github.com/triton-lang/triton-windows)." + ) + if q.dtype != v.dtype or k.dtype != v.dtype: + q = q.to(dtype=v.dtype) + k = k.to(dtype=v.dtype) + return na_attention_triton(q, k, v, list(attn.kernel_size)) + + +def warn_no_natten(*, backend: str) -> None: + """Emit the loud no-natten banner and which fallback backend was chosen.""" + logger.warning(_NO_NATTEN_WARNING) + logger.warning("DiffVAE NA fallback: using %s.", backend) + + +def fallback_na_attention() -> NAAttentionCallable: + """Pick Triton if usable, else eager; emit the no-natten warning.""" + if triton_na_available(): + warn_no_natten(backend="Triton na3d") + return TritonNaAttention() + warn_no_natten(backend="eager tiled SDPA na3d") + return EagerSdpaAttention() + + +__all__ = [ + "EagerSdpaAttention", + "TritonNaAttention", + "fallback_na_attention", + "na_attention_eager", + "na_attention_triton", + "triton_na_available", + "warn_no_natten", +] diff --git a/telefuser/models/ltx25/diff_vae/transformer/fallback_na/eager.py b/telefuser/models/ltx25/diff_vae/transformer/fallback_na/eager.py new file mode 100644 index 0000000..36ac07a --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/fallback_na/eager.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 Comfy Org. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Limited-workspace 3D neighborhood attention (NATTEN ``na3d`` semantics) in pure torch. +Vendored from comfy-kitchen ``backends/eager/na.py`` (Apache-2.0) for DiffVAE hosts +without natten or Triton. Queries are tiled; tiles that share window geometry stack +into batched ``scaled_dot_product_attention`` calls with one additive mask per group. +""" + +from __future__ import annotations + +import math + +import torch +from torch.nn import functional + +# Element budget for one tile's [Nq, Nk] attention mask (bounds the mask +# allocation and, on CPU, the math-backend score materialization). +NA_SCORE_BUDGET = 2**25 +# Element budget for the stacked K/V copies of one batched SDPA call on CUDA. +NA_KV_STACK_BUDGET = 2**28 + + +def _window_bounds(length: int, kernel: int, causal: bool) -> tuple[list[int], list[int]]: + """Per-index (start, end) of the attended window along one axis.""" + starts: list[int] = [] + ends: list[int] = [] + if causal: + for i in range(length): + starts.append(max(0, i - kernel + 1)) + ends.append(i + 1) + else: + kernel = min(kernel, length) + lo = length - kernel + half = kernel // 2 + for i in range(length): + start = min(max(i - half, 0), lo) + starts.append(start) + ends.append(start + kernel) + return starts, ends + + +def _pick_tiles(dims: tuple[int, int, int], kernels: list[int]) -> list[int]: + """Per-axis query-tile lengths keeping one tile's [Nq, Nk] under budget.""" + tiles = list(dims) + + def cost(ts: list[int]) -> int: + nq = math.prod(ts) + nk = math.prod(min(d, t + k - 1) for t, k, d in zip(ts, kernels, dims, strict=True)) + return nq * nk + + while cost(tiles) > NA_SCORE_BUDGET and max(tiles) > 1: + i = max(range(3), key=lambda a: tiles[a] / kernels[a]) + if tiles[i] <= 1: + break + tiles[i] = max(1, (tiles[i] + 1) // 2) + return tiles + + +def _group_mask( + rel_bounds: tuple[tuple[tuple[int, ...], tuple[int, ...]], ...], + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Additive ``[1, 1, Nq, Nk]`` mask for one tile-geometry group.""" + bools = [] + for starts, ends in rel_bounds: + st = torch.tensor(starts, device=device) + en = torch.tensor(ends, device=device) + kj = torch.arange(int(en.max()), device=device) + bools.append((kj[None, :] >= st[:, None]) & (kj[None, :] < en[:, None])) + visible = ( + bools[0][:, None, None, :, None, None] + & bools[1][None, :, None, None, :, None] + & bools[2][None, None, :, None, None, :] + ) + nq = visible.shape[0] * visible.shape[1] * visible.shape[2] + nk = visible.shape[3] * visible.shape[4] * visible.shape[5] + mask = torch.zeros((nq, nk), dtype=dtype, device=device) + mask.masked_fill_(~visible.reshape(nq, nk), torch.finfo(dtype).min) + return mask.reshape(1, 1, nq, nk) + + +def na3d( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kernel_size: list[int] | tuple[int, ...], + is_causal: list[bool] | None = None, + scale: float | None = None, +) -> torch.Tensor: + """3D neighborhood attention over ``(B, T, H, W, NH, HD)`` tensors. + ``scale`` defaults to ``head_dim**-0.5``. Pass ``scale=1.0`` when Q is already scaled. + """ + batch, t, h, w, nh, hd = q.shape + dims = (t, h, w) + causal = [False, False, False] if is_causal is None else list(is_causal) + kernels = [k_ if c else min(k_, d) for k_, c, d in zip(kernel_size, causal, dims, strict=True)] + if scale is None: + scale = hd**-0.5 + device = q.device + if scale != 1.0: + q = q * scale + + bounds = [_window_bounds(d, k_, c) for d, k_, c in zip(dims, kernels, causal, strict=True)] + tile_t, tile_h, tile_w = _pick_tiles(dims, [min(k_, d) for k_, d in zip(kernels, dims, strict=True)]) + + groups: dict[ + tuple[ + tuple[tuple[int, ...], tuple[int, ...]], + tuple[tuple[int, ...], tuple[int, ...]], + tuple[tuple[int, ...], tuple[int, ...]], + ], + list[tuple[tuple[slice, slice, slice], tuple[slice, slice, slice]]], + ] = {} + for t0 in range(0, t, tile_t): + t1 = min(t0 + tile_t, t) + rt0, rt1 = bounds[0][0][t0], bounds[0][1][t1 - 1] + rel_t = ( + tuple(s - rt0 for s in bounds[0][0][t0:t1]), + tuple(e - rt0 for e in bounds[0][1][t0:t1]), + ) + for h0 in range(0, h, tile_h): + h1 = min(h0 + tile_h, h) + rh0, rh1 = bounds[1][0][h0], bounds[1][1][h1 - 1] + rel_h = ( + tuple(s - rh0 for s in bounds[1][0][h0:h1]), + tuple(e - rh0 for e in bounds[1][1][h0:h1]), + ) + for w0 in range(0, w, tile_w): + w1 = min(w0 + tile_w, w) + rw0, rw1 = bounds[2][0][w0], bounds[2][1][w1 - 1] + rel_w = ( + tuple(s - rw0 for s in bounds[2][0][w0:w1]), + tuple(e - rw0 for e in bounds[2][1][w0:w1]), + ) + groups.setdefault((rel_t, rel_h, rel_w), []).append( + ( + (slice(t0, t1), slice(h0, h1), slice(w0, w1)), + (slice(rt0, rt1), slice(rh0, rh1), slice(rw0, rw1)), + ) + ) + + out = torch.empty((batch, t, h, w, nh, hd), device=device, dtype=v.dtype) + for rel, tiles in groups.items(): + mask = _group_mask(rel, q.dtype, device) + nq, nk = mask.shape[2], mask.shape[3] + g_max = max(1, NA_KV_STACK_BUDGET // max(1, batch * nh * nk * hd * 2)) if device.type == "cuda" else 1 + qs0, _ = tiles[0] + tq = qs0[0].stop - qs0[0].start + th = qs0[1].stop - qs0[1].start + tw = qs0[2].stop - qs0[2].start + for c0 in range(0, len(tiles), g_max): + chunk = tiles[c0 : c0 + g_max] + g = len(chunk) + q_s = torch.stack([q[:, qs[0], qs[1], qs[2]] for qs, _ in chunk]) + k_s = torch.stack([k[:, rs[0], rs[1], rs[2]] for _, rs in chunk]) + v_s = torch.stack([v[:, rs[0], rs[1], rs[2]] for _, rs in chunk]) + q_s = q_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * batch, nh, nq, hd) + k_s = k_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * batch, nh, nk, hd) + v_s = v_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * batch, nh, nk, hd) + o = functional.scaled_dot_product_attention(q_s, k_s, v_s, attn_mask=mask, scale=1.0) + o = o.view(g, batch, nh, tq, th, tw, hd).permute(0, 1, 3, 4, 5, 2, 6) + for i, (qs, _) in enumerate(chunk): + out[:, qs[0], qs[1], qs[2]] = o[i] + + return out diff --git a/telefuser/models/ltx25/diff_vae/transformer/fallback_na/triton_na.py b/telefuser/models/ltx25/diff_vae/transformer/fallback_na/triton_na.py new file mode 100644 index 0000000..e0e863f --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/fallback_na/triton_na.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 Comfy Org. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# ruff: noqa: ANN001, ANN202, PLR0912, PLR0913, PLR0915 +"""Triton 3D neighborhood attention (NATTEN ``na3d`` semantics). +Vendored from comfy-kitchen ``backends/triton/na.py`` (Apache-2.0) for DiffVAE hosts +without natten. One program handles a run of ``BLOCK_Q`` queries along W at a fixed +(t, h); online softmax, fp32 accumulation, no materialized scores. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +# Triton requires module-level JIT globals to be constexpr instances (not annotations). +_NEG_INF = tl.constexpr(-3.0e38) + + +@triton.jit +def _na3d_kernel( + q_ptr, + k_ptr, + v_ptr, + out_ptr, + t_size, + h_size, + w_size, + num_heads, + s_b, + s_t, + s_h, + s_w, + s_n, + scale, + kt: tl.constexpr, + kh: tl.constexpr, + kw: tl.constexpr, + causal_t: tl.constexpr, + causal_h: tl.constexpr, + causal_w: tl.constexpr, + hd: tl.constexpr, + hd_pad: tl.constexpr, + block_q: tl.constexpr, + block_k: tl.constexpr, + is_fp32: tl.constexpr, +): + pid_w = tl.program_id(0) + pid_th = tl.program_id(1) + pid_bn = tl.program_id(2) + + t_q = pid_th // h_size + h_q = pid_th % h_size + base = (pid_bn // num_heads) * s_b + (pid_bn % num_heads) * s_n + + w_off = pid_w * block_q + tl.arange(0, block_q) + w_valid = w_off < w_size + d_off = tl.arange(0, hd_pad) + d_mask = d_off < hd + + q_ptrs = q_ptr + base + t_q * s_t + h_q * s_h + w_off[:, None] * s_w + d_off[None, :] + q_blk = tl.load(q_ptrs, mask=w_valid[:, None] & d_mask[None, :], other=0.0) + + if causal_t: + t_lo = tl.maximum(t_q - kt + 1, 0) + t_hi = t_q + 1 + else: + t_lo = tl.minimum(tl.maximum(t_q - kt // 2, 0), t_size - kt) + t_hi = t_lo + kt + if causal_h: + h_lo = tl.maximum(h_q - kh + 1, 0) + h_hi = h_q + 1 + else: + h_lo = tl.minimum(tl.maximum(h_q - kh // 2, 0), h_size - kh) + h_hi = h_lo + kh + + w_q = tl.where(w_valid, w_off, w_size - 1) + if causal_w: + w_start = tl.maximum(w_q - kw + 1, 0) + w_end = w_q + 1 + blk_first = tl.minimum(pid_w * block_q, w_size - 1) + blk_last = tl.minimum(pid_w * block_q + block_q - 1, w_size - 1) + w_lo = tl.maximum(blk_first - kw + 1, 0) + w_hi = blk_last + 1 + else: + w_start = tl.minimum(tl.maximum(w_q - kw // 2, 0), w_size - kw) + w_end = w_start + kw + blk_first = tl.minimum(pid_w * block_q, w_size - 1) + blk_last = tl.minimum(pid_w * block_q + block_q - 1, w_size - 1) + w_lo = tl.minimum(tl.maximum(blk_first - kw // 2, 0), w_size - kw) + w_hi = tl.minimum(tl.maximum(blk_last - kw // 2, 0), w_size - kw) + kw + + m_i = tl.full((block_q,), _NEG_INF, dtype=tl.float32) + l_i = tl.zeros((block_q,), dtype=tl.float32) + acc = tl.zeros((block_q, hd_pad), dtype=tl.float32) + + for tk in range(t_lo, t_hi): + for hk in range(h_lo, h_hi): + plane = base + tk * s_t + hk * s_h + for wk0 in range(w_lo, w_hi, block_k): + wk = wk0 + tl.arange(0, block_k) + kmask = wk < w_hi + kv_ptrs = plane + wk[:, None] * s_w + d_off[None, :] + kv_mask = kmask[:, None] & d_mask[None, :] + k_blk = tl.load(k_ptr + kv_ptrs, mask=kv_mask, other=0.0) + if is_fp32: + s = tl.dot(q_blk, tl.trans(k_blk), input_precision="ieee") * scale + else: + s = tl.dot(q_blk, tl.trans(k_blk)) * scale + vis = (wk[None, :] >= w_start[:, None]) & (wk[None, :] < w_end[:, None]) & kmask[None, :] + s = tl.where(vis, s, _NEG_INF) + m_new = tl.maximum(m_i, tl.max(s, 1)) + alpha = tl.exp(m_i - m_new) + p = tl.exp(s - m_new[:, None]) + l_i = l_i * alpha + tl.sum(p, 1) + v_blk = tl.load(v_ptr + kv_ptrs, mask=kv_mask, other=0.0) + if is_fp32: + acc = acc * alpha[:, None] + tl.dot(p, v_blk, input_precision="ieee") + else: + acc = acc * alpha[:, None] + tl.dot(p.to(v_blk.dtype), v_blk) + m_i = m_new + + out = acc / tl.maximum(l_i, 1e-30)[:, None] + out_ptrs = out_ptr + base + t_q * s_t + h_q * s_h + w_off[:, None] * s_w + d_off[None, :] + tl.store(out_ptrs, out.to(out_ptr.dtype.element_ty), mask=w_valid[:, None] & d_mask[None, :]) + + +def na3d( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kernel_size: list[int] | tuple[int, ...], + is_causal: list[bool] | None = None, + scale: float | None = None, +) -> torch.Tensor: + """3D neighborhood attention over ``(B, T, H, W, NH, HD)`` tensors.""" + batch, t, h, w, nh, hd = q.shape + causal = [False, False, False] if is_causal is None else list(is_causal) + kt, kh, kw = (k_ if c else min(k_, d) for k_, c, d in zip(kernel_size, causal, (t, h, w), strict=True)) + if scale is None: + scale = hd**-0.5 + + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + out = torch.empty_like(q) + + hd_p = max(16, triton.next_power_of_2(hd)) + block_q = 16 + block_k = max(16, min(32, triton.next_power_of_2(min(w, block_q + kw)))) + + grid = (triton.cdiv(w, block_q), t * h, batch * nh) + _na3d_kernel[grid]( + q, + k, + v, + out, + t, + h, + w, + nh, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + q.stride(4), + scale, + kt=kt, + kh=kh, + kw=kw, + causal_t=causal[0], + causal_h=causal[1], + causal_w=causal[2], + hd=hd, + hd_pad=hd_p, + block_q=block_q, + block_k=block_k, + is_fp32=q.dtype == torch.float32, + num_warps=4, + ) + return out diff --git a/telefuser/models/ltx25/diff_vae/transformer/layers.py b/telefuser/models/ltx25/diff_vae/transformer/layers.py new file mode 100644 index 0000000..f41a1b3 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/layers.py @@ -0,0 +1,86 @@ +"""Shared small layers for the diffusion-VAE NA transformer stack.""" + +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F +from einops import rearrange +from torch import nn + + +class ChannelLinear(nn.Linear): + """``nn.Linear`` exposing ``in_channels``/``out_channels`` for config introspection.""" + + @property + def in_channels(self) -> int: + return self.in_features + + @property + def out_channels(self) -> int: + return self.out_features + + +class LinearPixelShuffleUpsample(nn.Module): + """Decoder-side resampler: Linear channel-expand, then channels-last PixelShuffle.""" + + def __init__( + self, + in_channels: int, + stride: tuple[int, int, int], + out_channels_reduction_factor: int = 1, + ) -> None: + super().__init__() + self.stride = stride + self.proj_out_channels = math.prod(stride) * in_channels // out_channels_reduction_factor + self.out_channels = self.proj_out_channels // math.prod(stride) + self.proj = nn.Linear(in_channels, self.proj_out_channels, bias=True) + + def forward(self, x: torch.Tensor, drop_leading_frame: bool = True) -> torch.Tensor: + """Upsample; when ``stride[0] == 2`` the pixel-shuffle produces a duplicate + leading frame that must be dropped to preserve the causal 1:2 (then + composed 1:8) frame mapping. ``drop_leading_frame`` gates that drop: it + must be ``True`` only for the chunk that contains the tensor's true + temporal origin (t=0). Tiled callers processing a later chunk in + isolation must pass ``False`` -- that chunk has no duplicate leading + frame of its own to drop, since the one duplicate frame in the full + (untiled) tensor belongs solely to the origin chunk. + """ + x = self.proj(x) + x = rearrange( + x, + "b t h w (c p1 p2 p3) -> b (t p1) (h p2) (w p3) c", + p1=self.stride[0], + p2=self.stride[1], + p3=self.stride[2], + ) + if self.stride[0] == 2 and drop_leading_frame: + x = x[:, 1:, :, :, :] + return x + + +class AdaLNZero(nn.Module): + """Per-block AdaLN-Zero modulation: ``t_emb`` -> 7 (scale/shift/gate) chunks. + Zero-init output projection so the block is an identity at every timestep + until the modulation pathway opens up during training. + """ + + NUM_CHUNKS: int = 7 # scale_msa, shift_msa, gate_msa, scale_mlp, shift_mlp, gate_mlp, gate_ctx + + def __init__(self, dim: int, t_emb_dim: int) -> None: + super().__init__() + self.dim = dim + self.proj = nn.Linear(t_emb_dim, self.NUM_CHUNKS * dim, bias=True) + nn.init.zeros_(self.proj.weight) + nn.init.zeros_(self.proj.bias) + + def forward(self, t_emb: torch.Tensor) -> tuple[torch.Tensor, ...]: + h = self.proj(F.silu(t_emb)) + chunks = h.chunk(self.NUM_CHUNKS, dim=-1) + return tuple(c[:, None, None, None, :] for c in chunks) + + +def modulate(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor) -> torch.Tensor: + """Apply AdaLN-style scale + shift modulation to a channels-last tensor.""" + return x * (1.0 + scale) + shift diff --git a/telefuser/models/ltx25/diff_vae/transformer/qkv.py b/telefuser/models/ltx25/diff_vae/transformer/qkv.py new file mode 100644 index 0000000..e285e81 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/qkv.py @@ -0,0 +1,34 @@ +"""Split Q/K/V projections for Neighborhood Attention.""" + +from __future__ import annotations + +import torch +from torch import nn + + +class QKVProjections(nn.Module): + """Three Q/K/V linears (separate GEMMs). + A fused ``nn.Linear(dim, 3*dim)`` + ``chunk`` looks cheaper, but on + ``(B,T,H,W,3*C)`` the chunks are non-contiguous (q|k|v interleaved per + spatial site). Making them contiguous while the fused buffer is alive + peaks at ``6*C`` channels (~+2 GB at stage-5) vs ``3*C`` from three + separate GEMMs -- and, unlike a fused tensor, Q/K/V here are independently + owned tensors with different lifetimes (e.g. Q is consumed immediately by + the score matmul while K/V live through the whole neighborhood gather), + so each can be freed as soon as its last use completes instead of all + three staying alive as long as the fused parent does. At production + shape the runtime difference is ~20 ms/step; the memory win dominates. + Checkpoint files still ship fused ``qkv.weight`` / ``qkv.bias``; those are + split into ``to_q`` / ``to_k`` / ``to_v`` during load by + ``DIFFUSION_VAE_DECODER_COMFY_KEYS_FILTER`` (``SDOps.with_kv_operation``). + """ + + def __init__(self, dim: int) -> None: + super().__init__() + self.dim = dim + self.to_q = nn.Linear(dim, dim, bias=True) + self.to_k = nn.Linear(dim, dim, bias=True) + self.to_v = nn.Linear(dim, dim, bias=True) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return self.to_q(x), self.to_k(x), self.to_v(x) diff --git a/telefuser/models/ltx25/diff_vae/transformer/rope.py b/telefuser/models/ltx25/diff_vae/transformer/rope.py new file mode 100644 index 0000000..ab6f1bd --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/rope.py @@ -0,0 +1,109 @@ +"""Shared absolute-RoPE helpers and knobs (no packaging policy). +Packaging lives with the consumer: + - det → ``det_attn_rope`` (opaque ``custom_op``) + - Combined diff → ``combined.attn`` (nested) + - Chunked diff → ``chunked.attn`` (in-slab) +""" + +from __future__ import annotations + +import torch +from torch import nn + +from telefuser.models.ltx25.diff_vae.transformer.rope_math import ( + DEFAULT_ABS_ROPE_NUM_TILES, + default_rope_dim_split, + h_positions, + rope_inv_freqs, + rot_abs_axis_impl, + t_positions, +) + +# Re-export math for existing call sites / tests. +__all__ = [ + "DEFAULT_ABS_ROPE_NUM_TILES", + "apply_abs_rope", + "apply_abs_rope_slab", + "configure_abs_rope", + "default_rope_dim_split", + "rope_inv_freqs", +] + + +def apply_abs_rope_slab( + x: torch.Tensor, + rope_split: tuple[int, int, int], + inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + *, + w_pos: torch.Tensor, + compute_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Rotate this W-extent (full volume or one slab). Returns a new tensor.""" + d_t, d_h, _ = rope_split + inv_t, inv_h, inv_w = inv_freqs + t = x.shape[1] + h = x.shape[2] + xt = rot_abs_axis_impl(x[..., :d_t], t_positions(t, x.device), inv_t, axis=1, compute_dtype=compute_dtype) + xh = rot_abs_axis_impl( + x[..., d_t : d_t + d_h], + h_positions(h, x.device), + inv_h, + axis=2, + compute_dtype=compute_dtype, + ) + xw = rot_abs_axis_impl(x[..., d_t + d_h :], w_pos, inv_w, axis=3, compute_dtype=compute_dtype) + return torch.cat([xt, xh, xw], dim=-1) + + +def apply_abs_rope( + x: torch.Tensor, + rope_split: tuple[int, int, int], + inv_freqs: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + *, + num_tiles: int = DEFAULT_ABS_ROPE_NUM_TILES, + compute_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Eager fixed-``num_tiles`` W-slab abs-RoPE (test / reference helper). + Production packaging is owned by ``det_attn_rope`` / ``diff_attn`` — not here. + """ + if num_tiles < 1: + raise ValueError(f"num_tiles must be >= 1, got {num_tiles}") + if compute_dtype not in (torch.float32, torch.bfloat16): + raise ValueError(f"compute_dtype must be float32 or bfloat16, got {compute_dtype}") + + slabs = torch.chunk(x, num_tiles, dim=3) + w_off = 0 + parts: list[torch.Tensor] = [] + for slab in slabs: + w_slab = slab.shape[3] + w_pos = torch.arange(w_slab, dtype=torch.float32, device=x.device) + w_off + parts.append( + apply_abs_rope_slab( + slab, + rope_split, + inv_freqs, + w_pos=w_pos, + compute_dtype=compute_dtype, + ) + ) + w_off = w_off + w_slab + return torch.cat(parts, dim=3) + + +def configure_abs_rope( + module_root: nn.Module, + *, + num_tiles: int = DEFAULT_ABS_ROPE_NUM_TILES, + compute_dtype: torch.dtype = torch.float32, +) -> None: + """Set abs-RoPE tile count / compute dtype on every ``NeighborhoodAttention3D``.""" + from telefuser.models.ltx25.diff_vae.transformer.attention import NeighborhoodAttention3D # noqa: PLC0415 + + if num_tiles < 1: + raise ValueError(f"num_tiles must be >= 1, got {num_tiles}") + if compute_dtype not in (torch.float32, torch.bfloat16): + raise ValueError(f"compute_dtype must be float32 or bfloat16, got {compute_dtype}") + for module in module_root.modules(): + if isinstance(module, NeighborhoodAttention3D): + module.rope_num_tiles = num_tiles + module.rope_compute_dtype = compute_dtype diff --git a/telefuser/models/ltx25/diff_vae/transformer/rope_math.py b/telefuser/models/ltx25/diff_vae/transformer/rope_math.py new file mode 100644 index 0000000..a671a8c --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/rope_math.py @@ -0,0 +1,61 @@ +"""Shared absolute-RoPE math helpers (no consumer policy).""" + +from __future__ import annotations + +import numpy as np +import torch + +DEFAULT_ABS_ROPE_NUM_TILES = 4 + + +def t_positions(t: int, device: torch.device) -> torch.Tensor: + return torch.arange(t, dtype=torch.float32, device=device) + + +def h_positions(h: int, device: torch.device) -> torch.Tensor: + return torch.arange(h, dtype=torch.float32, device=device) + + +def default_rope_dim_split(head_dim: int) -> tuple[int, int, int]: + """Default split of head_dim across (T, H, W) RoPE chunks.""" + assert head_dim % 8 == 0, f"head_dim={head_dim} must be a multiple of 8 for default split" + d_t = (head_dim // 4) // 2 * 2 + d_hw = (head_dim - d_t) // 2 + if d_hw % 2 != 0: + d_t -= 2 + d_hw = (head_dim - d_t) // 2 + assert d_t > 0 + assert d_hw > 0 + return (d_t, d_hw, d_hw) + + +def rope_inv_freqs(dim: int, base: float = 10000.0) -> torch.Tensor: + """Inverse RoPE frequencies: ``1 / base**(i/dim)`` for ``i`` in ``[0, dim, 2)``.""" + assert dim % 2 == 0, f"RoPE dim must be even, got {dim}" + exponents = np.arange(0, dim, 2, dtype=np.float64) / dim + inv_freqs = 1.0 / np.power(float(base), exponents) + return torch.from_numpy(inv_freqs).to(torch.float32) + + +def rot_abs_axis_impl( + xc: torch.Tensor, + pos: torch.Tensor, + inv: torch.Tensor, + axis: int, + *, + compute_dtype: torch.dtype, +) -> torch.Tensor: + """Absolute RoPE on one axis chunk ``xc[..., D]`` (D even) → new tensor.""" + out_dtype = xc.dtype + pairs = xc.reshape(*xc.shape[:-1], xc.shape[-1] // 2, 2) + xe = pairs[..., 0].to(compute_dtype) + xo = pairs[..., 1].to(compute_dtype) + shape = [1, 1, 1, 1, 1, inv.shape[0]] + shape[axis] = pos.shape[0] + ang = (pos[:, None] * inv[None, :]).reshape(shape) + c = ang.cos().to(compute_dtype) + s = ang.sin().to(compute_dtype) + re = xe * c - xo * s + ro = xe * s + xo * c + out = torch.stack([re, ro], dim=-1).reshape(xc.shape) + return out.to(out_dtype) if out.dtype != out_dtype else out diff --git a/telefuser/models/ltx25/diff_vae/transformer/swiglu.py b/telefuser/models/ltx25/diff_vae/transformer/swiglu.py new file mode 100644 index 0000000..f03f382 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/transformer/swiglu.py @@ -0,0 +1,479 @@ +# ruff: noqa: ANN001, N803, ANN202, PLR0913 + +"""SwiGLU shell, tiled kernel, det plain residual, and shared tile/fuse helpers. +Pathway residuals: + - Combined → ``combined.mlp.residual_mlp`` + - Chunked → ``chunked.mlp.residual_modulating_mlp`` +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Final + +import torch +import torch.nn.functional as F +from torch import nn + +from telefuser.models.ltx25.diff_vae.tiling import DimensionInterval, split_by_count, split_by_size + +logger = logging.getLogger(__name__) + +try: + import triton + import triton.language as tl + + _TRITON_AVAILABLE = True +except ImportError: # pragma: no cover + triton = None # type: ignore[assignment] + tl = None # type: ignore[assignment] + _TRITON_AVAILABLE = False + +# Match visual_fidelity ``DEFAULT_CHUNK_TOKENS`` — hard token cap for peak VRAM. +DEFAULT_SWIGLU_TILE_SIZE: Final[int] = 16_384 +DEFAULT_SWIGLU_TILES: Final[int] = 4 + + +@dataclass(frozen=True) +class SwiGLUTileSpec: + """Token tiling: provide exactly one of ``num_tiles`` or ``tile_size``.""" + + num_tiles: int | None = None + tile_size: int | None = None + + def __post_init__(self) -> None: + if (self.num_tiles is None) == (self.tile_size is None): + raise ValueError("Provide exactly one of num_tiles or tile_size") + if self.num_tiles is not None and self.num_tiles < 1: + raise ValueError(f"num_tiles must be >= 1, got {self.num_tiles}") + if self.tile_size is not None and self.tile_size < 1: + raise ValueError(f"tile_size must be >= 1, got {self.tile_size}") + + @classmethod + def by_count(cls, num_tiles: int = DEFAULT_SWIGLU_TILES) -> SwiGLUTileSpec: + return cls(num_tiles=num_tiles) + + @classmethod + def by_size(cls, tile_size: int = DEFAULT_SWIGLU_TILE_SIZE) -> SwiGLUTileSpec: + return cls(tile_size=tile_size) + + +DEFAULT_SWIGLU_TILE_SPEC = SwiGLUTileSpec(tile_size=DEFAULT_SWIGLU_TILE_SIZE) + + +def triton_swiglu_available() -> bool: + """True when a CUDA-capable Triton install can run the fused up-mul kernel.""" + return _TRITON_AVAILABLE and torch.cuda.is_available() + + +def token_intervals(n_tok: int, tile: SwiGLUTileSpec) -> list[DimensionInterval]: + """Split ``n_tok`` tokens per ``tile`` (count or size). Overlap is always 0.""" + if n_tok < 0: + raise ValueError(f"n_tok must be >= 0, got {n_tok}") + if n_tok == 0: + return [] + if tile.num_tiles is not None: + tiles = min(tile.num_tiles, n_tok) + if tiles <= 1: + return [DimensionInterval(start=0, end=n_tok, left_ramp=0, right_ramp=0)] + return list(split_by_count(num_tiles=tiles, overlap=0)(n_tok).intervals) + assert tile.tile_size is not None + if n_tok <= tile.tile_size: + return [DimensionInterval(start=0, end=n_tok, left_ramp=0, right_ramp=0)] + return list(split_by_size(tile.tile_size, overlap=0)(n_tok).intervals) + + +if _TRITON_AVAILABLE: + + @triton.jit + def _fused_up_mul_kernel( + x_ptr, + w_ptr, + g_ptr, + M, + K, + N, + stride_xm, + stride_xk, + stride_wn, + stride_wk, + stride_gm, + stride_gn, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + """``g = g * (x @ w.T)`` with ``w`` shaped ``(N, K)`` (nn.Linear layout).""" + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + mask_m = offs_m < M + mask_n = offs_n < N + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k0 in range(0, K, BLOCK_K): + offs_k = k0 + tl.arange(0, BLOCK_K) + mask_k = offs_k < K + x = tl.load( + x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk, + mask=mask_m[:, None] & mask_k[None, :], + other=0.0, + ) + w = tl.load( + w_ptr + offs_n[None, :] * stride_wn + offs_k[:, None] * stride_wk, + mask=mask_k[:, None] & mask_n[None, :], + other=0.0, + ) + acc += tl.dot(x, w) + + g = tl.load( + g_ptr + offs_m[:, None] * stride_gm + offs_n[None, :] * stride_gn, + mask=mask_m[:, None] & mask_n[None, :], + other=0.0, + ) + out = g * acc.to(g.dtype) + tl.store( + g_ptr + offs_m[:, None] * stride_gm + offs_n[None, :] * stride_gn, + out, + mask=mask_m[:, None] & mask_n[None, :], + ) + + @triton.jit + def _fused_gate_up_swiglu_kernel( + x_ptr, + w_gate_ptr, + w_up_ptr, + out_ptr, + M, + K, + N, + stride_xm, + stride_xk, + stride_gate_n, + stride_gate_k, + stride_up_n, + stride_up_k, + stride_om, + stride_on, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + """One x load → gate+up GEMMs → SiLU → gated product into ``out`` (BF16). + Rounding matches the unfused BF16 path: SiLU is rounded to BF16 before the + product, then ``(silu_f32 * up_f32).to(bf16)`` for the store. + """ + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + mask_m = offs_m < M + mask_n = offs_n < N + + gate_acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + up_acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k0 in range(0, K, BLOCK_K): + offs_k = k0 + tl.arange(0, BLOCK_K) + mask_k = offs_k < K + x = tl.load( + x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk, + mask=mask_m[:, None] & mask_k[None, :], + other=0.0, + ) + gate_w = tl.load( + w_gate_ptr + offs_n[None, :] * stride_gate_n + offs_k[:, None] * stride_gate_k, + mask=mask_k[:, None] & mask_n[None, :], + other=0.0, + ) + up_w = tl.load( + w_up_ptr + offs_n[None, :] * stride_up_n + offs_k[:, None] * stride_up_k, + mask=mask_k[:, None] & mask_n[None, :], + other=0.0, + ) + gate_acc += tl.dot(x, gate_w) + up_acc += tl.dot(x, up_w) + + silu_bf16 = (gate_acc * tl.sigmoid(gate_acc)).to(tl.bfloat16) + product = (silu_bf16.to(tl.float32) * up_acc).to(tl.bfloat16) + tl.store( + out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on, + product, + mask=mask_m[:, None] & mask_n[None, :], + ) + + +_DUAL_GATE_UP_DIM_MIN: Final[int] = 256 +_DUAL_GATE_UP_DIM_MAX: Final[int] = 2048 + + +def dual_gate_up_triton_eligible( + x: torch.Tensor, + w_gate: torch.Tensor, + w_up: torch.Tensor, +) -> bool: + """True when the fused gate+up SwiGLU Triton kernel can run. + Requires BF16 activations/weights, ``hidden == 4 * dim``, and ``dim`` in + ``[256, 2048]`` (DiffVAE MLP shapes). + """ + if not triton_swiglu_available() or not x.is_cuda: + return False + if x.dtype != torch.bfloat16 or w_gate.dtype != torch.bfloat16 or w_up.dtype != torch.bfloat16: + return False + dim = x.shape[-1] + hidden = w_gate.shape[0] + if w_up.shape[0] != hidden or w_gate.shape[1] != dim or w_up.shape[1] != dim: + return False + if hidden != 4 * dim: + return False + return _DUAL_GATE_UP_DIM_MIN <= dim <= _DUAL_GATE_UP_DIM_MAX + + +def _fused_gate_up_swiglu_triton( + x: torch.Tensor, + w_gate: torch.Tensor, + w_up: torch.Tensor, + out: torch.Tensor, +) -> None: + """``out = silu(x @ W_gateᵀ) * (x @ W_upᵀ)`` with one x-tile load (BF16).""" + if not _TRITON_AVAILABLE: + raise RuntimeError("Triton is not available") + m, k = x.shape + n = w_gate.shape[0] + block_m, block_n, block_k = 64, 64, 32 + grid = (triton.cdiv(m, block_m), triton.cdiv(n, block_n)) + _fused_gate_up_swiglu_kernel[grid]( + x, + w_gate, + w_up, + out, + m, + k, + n, + x.stride(0), + x.stride(1), + w_gate.stride(0), + w_gate.stride(1), + w_up.stride(0), + w_up.stride(1), + out.stride(0), + out.stride(1), + BLOCK_M=block_m, + BLOCK_N=block_n, + BLOCK_K=block_k, + ) + + +def _fused_up_mul_triton(x: torch.Tensor, w_up: torch.Tensor, silu_gate: torch.Tensor) -> None: + """Inplace: ``silu_gate *= (x @ w_up.T)``.""" + if not _TRITON_AVAILABLE: + raise RuntimeError("Triton is not available") + m, k = x.shape + n = w_up.shape[0] + block_m, block_n, block_k = 64, 64, 32 + grid = (triton.cdiv(m, block_m), triton.cdiv(n, block_n)) + _fused_up_mul_kernel[grid]( + x, + w_up, + silu_gate, + m, + k, + n, + x.stride(0), + x.stride(1), + w_up.stride(0), + w_up.stride(1), + silu_gate.stride(0), + silu_gate.stride(1), + BLOCK_M=block_m, + BLOCK_N=block_n, + BLOCK_K=block_k, + ) + + +def _fused_up_mul_torch(x: torch.Tensor, w_up: torch.Tensor, silu_gate: torch.Tensor) -> None: + """PyTorch fallback: temporary ``up`` chunk, then inplace mul into ``silu_gate``.""" + up = F.linear(x, w_up) + silu_gate.mul_(up) + + +def swiglu( + x: torch.Tensor, + w_gate: torch.Tensor, + w_up: torch.Tensor, + w_down: torch.Tensor, +) -> torch.Tensor: + """Reference: ``w_down(silu(x@W_gateᵀ) * (x@W_upᵀ))`` (full materialization).""" + return F.linear(F.silu(F.linear(x, w_gate)) * F.linear(x, w_up), w_down) + + +def _tile_from_op_args(tile_size: int, num_tiles: int) -> SwiGLUTileSpec: + """Decode custom_op int knobs: exactly one of ``tile_size`` / ``num_tiles`` is > 0.""" + if (tile_size > 0) == (num_tiles > 0): + raise ValueError(f"need exactly one of tile_size>0 or num_tiles>0, got {tile_size=}, {num_tiles=}") + if tile_size > 0: + return SwiGLUTileSpec(tile_size=tile_size) + return SwiGLUTileSpec(num_tiles=num_tiles) + + +def _swiglu_chunked_impl( + x: torch.Tensor, + w_gate: torch.Tensor, + w_up: torch.Tensor, + w_down: torch.Tensor, + intervals: list[DimensionInterval], + *, + use_triton: bool, +) -> torch.Tensor: + """Chunked SwiGLU: one reusable ``(chunk, hidden)`` workspace (runs via custom_op).""" + if use_triton and not triton_swiglu_available(): + raise RuntimeError("use_triton=True but Triton/CUDA is unavailable") + + # RMSNorm under bf16 autocast can leave activations in float32 while Linear weights + # stay bf16. Inside this custom_op autocast does not cast for ``torch.mm``. + if x.dtype != w_gate.dtype: + x = x.to(dtype=w_gate.dtype) + + leading = x.shape[:-1] + dim = x.shape[-1] + x_flat = x.reshape(-1, dim).contiguous() + n_tok = x_flat.shape[0] + hidden = w_gate.shape[0] + if w_up.shape[0] != hidden or w_gate.shape[1] != dim or w_up.shape[1] != dim: + raise ValueError( + f"weight shapes incompatible with x[..., {dim}]: gate={tuple(w_gate.shape)} up={tuple(w_up.shape)}" + ) + if w_down.shape[1] != hidden or w_down.shape[0] != dim: + raise ValueError(f"w_down {tuple(w_down.shape)} incompatible with hidden={hidden} dim={dim}") + if n_tok == 0: + return x + + out_flat = torch.empty((n_tok, dim), device=x.device, dtype=x.dtype) + max_chunk = max((iv.end - iv.start for iv in intervals), default=0) + workspace = torch.empty((max_chunk, hidden), device=x.device, dtype=x.dtype) + use_dual = bool(use_triton) and dual_gate_up_triton_eligible(x_flat, w_gate, w_up) + w_gate_c = w_gate.contiguous() if use_dual else w_gate + w_up_c = w_up.contiguous() if use_triton else w_up + + for iv in intervals: + start, end = iv.start, iv.end + if start >= end: + continue + xc = x_flat[start:end] + ws = workspace[: end - start] + if use_dual: + _fused_gate_up_swiglu_triton(xc, w_gate_c, w_up_c, ws) + else: + torch.mm(xc, w_gate.t(), out=ws) + F.silu(ws, inplace=True) + if use_triton: + _fused_up_mul_triton(xc, w_up_c, ws) + else: + _fused_up_mul_torch(xc, w_up, ws) + torch.mm(ws, w_down.t(), out=out_flat[start:end]) + + return out_flat.view(*leading, dim) + + +@torch.library.custom_op("ltx_core::swiglu_tiled", mutates_args=()) +def _swiglu_tiled_op( + x: torch.Tensor, + w_gate: torch.Tensor, + w_up: torch.Tensor, + w_down: torch.Tensor, + tile_size: int, + num_tiles: int, + use_triton: bool, +) -> torch.Tensor: + """Opaque op: memory-efficient chunked SwiGLU (compile-friendly; no Dynamo unroll).""" + tile = _tile_from_op_args(tile_size, num_tiles) + n_tok = x.reshape(-1, x.shape[-1]).shape[0] + return _swiglu_chunked_impl(x, w_gate, w_up, w_down, token_intervals(n_tok, tile), use_triton=use_triton) + + +@_swiglu_tiled_op.register_fake +def _swiglu_tiled_fake( + x: torch.Tensor, + w_gate: torch.Tensor, + w_up: torch.Tensor, + w_down: torch.Tensor, + tile_size: int, + num_tiles: int, + use_triton: bool, +) -> torch.Tensor: + del w_gate, w_up, w_down, tile_size, num_tiles, use_triton + return torch.empty(x.shape, device=x.device, dtype=x.dtype) + + +def swiglu_tiled( + x: torch.Tensor, + w_gate: torch.Tensor, + w_up: torch.Tensor, + w_down: torch.Tensor, + tile: SwiGLUTileSpec, + *, + use_triton: bool | None = None, +) -> torch.Tensor: + """Memory-efficient SwiGLU on the full activation; tiling is internal (custom_op).""" + tile_size = int(tile.tile_size) if tile.tile_size is not None else 0 + num_tiles = int(tile.num_tiles) if tile.num_tiles is not None else 0 + if use_triton is None: + use_triton = triton_swiglu_available() and x.is_cuda + return _swiglu_tiled_op(x, w_gate, w_up, w_down, tile_size, num_tiles, bool(use_triton)) + + +def plain_mlp(x: torch.Tensor, mlp: nn.Module, norm: nn.RMSNorm, tile: SwiGLUTileSpec) -> torch.Tensor: + """Det ``NABlock``: ``x + swiglu_tiled(norm(x))`` (no AdaLN).""" + y = norm(x) + if y.numel() == 0: + return x + return x + swiglu_tiled(y, mlp.w_gate.weight, mlp.w_up.weight, mlp.w_down.weight, tile) + + +class SwiGLU(nn.Module): + """Gated MLP weights: ``w_down(silu(w_gate(x)) * w_up(x))``.""" + + def __init__( + self, + dim: int, + hidden_dim: int, + tile: SwiGLUTileSpec = DEFAULT_SWIGLU_TILE_SPEC, + ) -> None: + super().__init__() + self.w_up = nn.Linear(dim, hidden_dim, bias=False) + self.w_gate = nn.Linear(dim, hidden_dim, bias=False) + self.w_down = nn.Linear(hidden_dim, dim, bias=False) + self.tile = tile + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Tiled SwiGLU on ``x`` (tests / direct calls). Block residuals use pathway fns.""" + if x.numel() == 0: + return x + return swiglu_tiled( + x, + self.w_gate.weight, + self.w_up.weight, + self.w_down.weight, + self.tile, + ) + + +def configure_swiglu_tile( + module_root: nn.Module, + *, + num_tiles: int | None = None, + tile_size: int | None = None, +) -> None: + """Set tile spec on every ``SwiGLU`` under ``module_root``. + Pass exactly one of ``num_tiles`` or ``tile_size``. Calling with both omitted + is a no-op (leaves each module's existing ``tile``). + """ + if num_tiles is None and tile_size is None: + return + tile = SwiGLUTileSpec(num_tiles=num_tiles, tile_size=tile_size) + for module in module_root.modules(): + if isinstance(module, SwiGLU): + module.tile = tile diff --git a/telefuser/models/ltx25/diff_vae/types.py b/telefuser/models/ltx25/diff_vae/types.py new file mode 100644 index 0000000..115b1e3 --- /dev/null +++ b/telefuser/models/ltx25/diff_vae/types.py @@ -0,0 +1,300 @@ +from dataclasses import dataclass, replace +from typing import NamedTuple + +import torch + + +class VideoPixelShape(NamedTuple): + """ + Shape of the tensor representing the video pixel array. Assumes BGR channel format. + """ + + batch: int + frames: int + height: int + width: int + fps: float + + +class SpatioTemporalScaleFactors(NamedTuple): + """ + Describes the spatiotemporal downscaling between decoded video space and + the corresponding VAE latent grid. + Field order matches the (frame/time, height, width) axis layout used by + latent tensors and meshgrid coordinates elsewhere in the codebase. + """ + + time: int + height: int + width: int + + @classmethod + def default(cls) -> "SpatioTemporalScaleFactors": + return cls(time=8, height=32, width=32) + + @classmethod + def from_blocks(cls, blocks: list, patch_size: int) -> "SpatioTemporalScaleFactors": + """Derive the scale factors from a VAE encoder/decoder block list. + Each ``compress_*`` block halves (encoder) or doubles (decoder) its target + axes by a stride of 2, independent of any channel ``multiplier``. The initial + patchify contributes an extra ``patch_size`` of spatial compression. Deriving + the factors from the blocks keeps a single source of truth that stays correct + across VAE variants (e.g. the 32x32x8 default and the 16x16x4 variant) instead + of relying on a hardcoded constant. + """ + spatial_steps = 0 + temporal_steps = 0 + for block_name, _ in blocks: + if block_name.startswith(("compress_space", "compress_all")): + spatial_steps += 1 + if block_name.startswith(("compress_time", "compress_all")): + temporal_steps += 1 + spatial = patch_size * (2**spatial_steps) + return cls(time=2**temporal_steps, height=spatial, width=spatial) + + @classmethod + def from_model_config(cls, model_config: dict) -> "SpatioTemporalScaleFactors": + """Derive the video scale factors from a checkpoint's model config dict. + Reads the embedded VAE block list (see ``from_blocks``). Falls back to the + default when the config carries no VAE block list -- either no ``vae`` section + or a ``vae`` section without encoder/decoder blocks (e.g. audio-only + checkpoints), where video tools are never built. + """ + vae_config = model_config.get("vae", {}) + blocks = vae_config.get("encoder_blocks") or vae_config.get("decoder_blocks") + if not blocks: + return cls.default() + return cls.from_blocks(blocks, vae_config.get("patch_size", 4)) + + +VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default() + + +class VideoLatentShape(NamedTuple): + """ + Shape of the tensor representing video in VAE latent space. + The latent representation is a 5D tensor with dimensions ordered as + (batch, channels, frames, height, width). Spatial and temporal dimensions + are downscaled relative to pixel space according to the VAE's scale factors. + """ + + batch: int + channels: int + frames: int + height: int + width: int + + def to_torch_shape(self) -> torch.Size: + return torch.Size([self.batch, self.channels, self.frames, self.height, self.width]) + + @staticmethod + def from_torch_shape(shape: torch.Size) -> "VideoLatentShape": + return VideoLatentShape( + batch=shape[0], + channels=shape[1], + frames=shape[2], + height=shape[3], + width=shape[4], + ) + + def token_count(self) -> int: + """Number of tokens after patchification with the default patch size of 1.""" + return self.frames * self.height * self.width + + def mask_shape(self) -> "VideoLatentShape": + return self._replace(channels=1) + + @staticmethod + def from_pixel_shape( + shape: VideoPixelShape, + latent_channels: int = 128, + scale_factors: SpatioTemporalScaleFactors = VIDEO_SCALE_FACTORS, + ) -> "VideoLatentShape": + frames = (shape.frames - 1) // scale_factors.time + 1 + height = shape.height // scale_factors.height + width = shape.width // scale_factors.width + + return VideoLatentShape( + batch=shape.batch, + channels=latent_channels, + frames=frames, + height=height, + width=width, + ) + + def upscale(self, scale_factors: SpatioTemporalScaleFactors = VIDEO_SCALE_FACTORS) -> "VideoLatentShape": + return self._replace( + channels=3, + frames=(self.frames - 1) * scale_factors.time + 1, + height=self.height * scale_factors.height, + width=self.width * scale_factors.width, + ) + + +class AudioLatentShape(NamedTuple): + """ + Shape of audio in VAE latent space: (batch, channels, frames, mel_bins). + mel_bins is the number of frequency bins from the mel-spectrogram encoding. + """ + + batch: int + channels: int + frames: int + mel_bins: int + + def to_torch_shape(self) -> torch.Size: + return torch.Size([self.batch, self.channels, self.frames, self.mel_bins]) + + def token_count(self) -> int: + """Number of tokens after patchification.""" + return self.frames + + def mask_shape(self) -> "AudioLatentShape": + return self._replace(channels=1, mel_bins=1) + + @staticmethod + def from_torch_shape(shape: torch.Size) -> "AudioLatentShape": + return AudioLatentShape( + batch=shape[0], + channels=shape[1], + frames=shape[2], + mel_bins=shape[3], + ) + + @staticmethod + def from_duration( + batch: int, + duration: float, + channels: int = 8, + mel_bins: int = 16, + sample_rate: int = 16000, + hop_length: int = 160, + audio_latent_downsample_factor: int = 4, + ) -> "AudioLatentShape": + latents_per_second = float(sample_rate) / float(hop_length) / float(audio_latent_downsample_factor) + + return AudioLatentShape( + batch=batch, + channels=channels, + frames=round(duration * latents_per_second), + mel_bins=mel_bins, + ) + + @staticmethod + def from_video_pixel_shape( + shape: VideoPixelShape, + channels: int = 8, + mel_bins: int = 16, + sample_rate: int = 16000, + hop_length: int = 160, + audio_latent_downsample_factor: int = 4, + ) -> "AudioLatentShape": + return AudioLatentShape.from_duration( + batch=shape.batch, + duration=float(shape.frames) / float(shape.fps), + channels=channels, + mel_bins=mel_bins, + sample_rate=sample_rate, + hop_length=hop_length, + audio_latent_downsample_factor=audio_latent_downsample_factor, + ) + + +@dataclass(frozen=True) +class Audio: + """ + Container for decoded audio samples and metadata. + Attributes: + waveform: Audio waveform tensor. + sampling_rate: Sampling rate (Hz) of the waveform. + """ + + waveform: torch.Tensor + sampling_rate: int + + def to(self, **kwargs: object) -> "Audio": + return replace(self, waveform=self.waveform.to(**kwargs)) + + +@dataclass(frozen=True) +class GeneratedKeyframeLayout: + """Where a state's generated-keyframe slot tokens live, and what they represent. + Recorded by :class:`~ltx_core.conditioning.types.keyframe_slots.VideoGeneratedKeyframeSlots` + when it appends the slots, so they can later be located and extracted *exactly* rather + than by assuming they are the trailing tokens. Conditioning items are applied in list + order and each appends to the end, so a state built with slots plus any other appending + conditioning item has no fixed trailing layout. + Attributes: + pixel_frame_indices: Target pixel-frame index of each slot, in token order. + tokens_per_keyframe: Number of tokens one slot occupies (one latent frame's worth). + first_token: Index of the first slot token in the token sequence. + """ + + pixel_frame_indices: tuple[int, ...] + tokens_per_keyframe: int + first_token: int + + @property + def num_keyframes(self) -> int: + return len(self.pixel_frame_indices) + + @property + def num_tokens(self) -> int: + return self.num_keyframes * self.tokens_per_keyframe + + @property + def token_slice(self) -> slice: + return slice(self.first_token, self.first_token + self.num_tokens) + + +@dataclass(frozen=True) +class LatentState: + """ + State of latents during the diffusion denoising process. + Attributes: + latent: The current noisy latent tensor being denoised. + denoise_mask: Mask encoding the denoising strength for each token (1 = full denoising, 0 = no denoising). + positions: Positional indices for each latent element, used for positional embeddings. + clean_latent: Initial state of the latent before denoising, may include conditioning latents. + attention_mask: Optional 2D self-attention mask of shape (B, T, T). Values in [0, 1] where 1 = full attention, + 0 = no attention. None means full attention everywhere. Built incrementally by conditioning items. + keyframes_mask: Optional per-token marker of shape (B, T, 1) -- same layout as + ``denoise_mask`` -- non-zero on tokens whose latent encodes a *single standalone pixel + frame* rather than the usual multi-frame span. That set is the target's first latent + frame (the video encoder is causal, so its first temporal latent frame covers 1 pixel + frame while the rest cover 8) plus any generated keyframe slots. Selects the tokens + that receive the model's learned keyframe absolute-position embedding; ignored + entirely by models built without ``use_keyframes_abs_pos_embedding``. + generated_keyframe_layout: Set when generated keyframe slots were appended; locates them. + generated_keyframes: Populated by ``clear_conditioning`` when a layout is present: the + denoised slot content as an unpatchified ``(B, C, K, H, W)`` latent, one latent frame + per keyframe. Each frame must be decoded as a standalone one-frame clip, never as a + K-frame video -- a causal decode would blend slots that were never adjacent. + frozen: When True, this stream is held fixed: token denoising is disabled (``denoise_mask`` + should be all zeros; pipeline builders enforce that) and the scalar noise level used for + prompt / cross-modality AdaLN gates is forced to 0 when the state is converted for the + transformer. + """ + + latent: torch.Tensor + denoise_mask: torch.Tensor + positions: torch.Tensor + clean_latent: torch.Tensor + attention_mask: torch.Tensor | None = None + keyframes_mask: torch.Tensor | None = None + generated_keyframe_layout: GeneratedKeyframeLayout | None = None + generated_keyframes: torch.Tensor | None = None + frozen: bool = False + + def clone(self) -> "LatentState": + return LatentState( + latent=self.latent.clone(), + denoise_mask=self.denoise_mask.clone(), + positions=self.positions.clone(), + clean_latent=self.clean_latent.clone(), + attention_mask=self.attention_mask.clone() if self.attention_mask is not None else None, + keyframes_mask=self.keyframes_mask.clone() if self.keyframes_mask is not None else None, + generated_keyframe_layout=self.generated_keyframe_layout, + generated_keyframes=(self.generated_keyframes.clone() if self.generated_keyframes is not None else None), + frozen=self.frozen, + ) diff --git a/telefuser/models/ltx25/duration.py b/telefuser/models/ltx25/duration.py new file mode 100644 index 0000000..16adde0 --- /dev/null +++ b/telefuser/models/ltx25/duration.py @@ -0,0 +1,142 @@ +"""LTX-2.5 DurationHead checkpoint loader and frame-grid resolution.""" + +from __future__ import annotations + +from pathlib import Path + +import torch +from safetensors import safe_open +from torch import nn + +from .checkpoint import inspect_checkpoint + + +class LTX25AttentionPooler(nn.Module): + """Learned-query cross-attention pooler used by the DurationHead.""" + + def __init__(self, hidden_dim: int, num_queries: int, num_heads: int) -> None: + super().__init__() + self.query_tokens = nn.Parameter(torch.randn(num_queries, hidden_dim) * 0.02) + self.cross_attn = nn.MultiheadAttention(hidden_dim, num_heads, batch_first=True) + + def forward(self, tokens: torch.Tensor) -> torch.Tensor: + queries = self.query_tokens.unsqueeze(0).expand(tokens.shape[0], -1, -1) + pooled, _ = self.cross_attn(queries, tokens, tokens, need_weights=False) + return pooled + + +class LTX25DurationHead(nn.Module): + """Predict a duration in seconds from LTX-2.5 connector output tokens.""" + + def __init__( + self, + *, + video_cross_attention_dim: int, + audio_cross_attention_dim: int, + pooler_hidden_dim: int = 256, + num_queries: int = 1, + num_pooler_heads: int = 4, + mlp_hidden: int = 256, + ) -> None: + super().__init__() + self.video_input_proj = nn.Linear(video_cross_attention_dim, pooler_hidden_dim) + self.video_modality_emb = nn.Parameter(torch.randn(pooler_hidden_dim) * 0.02) + self.audio_input_proj = nn.Linear(audio_cross_attention_dim, pooler_hidden_dim) + self.audio_modality_emb = nn.Parameter(torch.randn(pooler_hidden_dim) * 0.02) + self.attention_pooler = LTX25AttentionPooler(pooler_hidden_dim, num_queries, num_pooler_heads) + self.mlp_hidden = nn.Linear(pooler_hidden_dim * num_queries, mlp_hidden) + self.mlp_out = nn.Linear(mlp_hidden, 1) + + def forward( + self, + video_tokens: torch.Tensor | None = None, + audio_tokens: torch.Tensor | None = None, + ) -> torch.Tensor: + """Return a positive duration prediction in seconds for every batch item.""" + if video_tokens is None and audio_tokens is None: + raise ValueError("LTX25DurationHead requires video_tokens or audio_tokens") + groups: list[torch.Tensor] = [] + if video_tokens is not None: + groups.append(self.video_input_proj(video_tokens) + self.video_modality_emb) + if audio_tokens is not None: + groups.append(self.audio_input_proj(audio_tokens) + self.audio_modality_emb) + pooled = self.attention_pooler(torch.cat(groups, dim=1)).flatten(1) + hidden = torch.nn.functional.gelu(self.mlp_hidden(pooled), approximate="tanh") + return self.mlp_out(hidden).squeeze(-1).exp() + + @classmethod + def from_checkpoint( + cls, + checkpoint_path: str | Path, + *, + device: torch.device | str = "cpu", + torch_dtype: torch.dtype = torch.bfloat16, + ) -> "LTX25DurationHead": + """Construct and strictly load the standalone LTX-2.5 DurationHead.""" + metadata = inspect_checkpoint(checkpoint_path) + transformer = metadata.config.get("transformer", {}) + duration = metadata.config.get("duration_head", {}) + if not isinstance(transformer, dict) or not isinstance(duration, dict): + raise ValueError("LTX-2.5 DurationHead metadata must contain transformer and duration_head objects") + with torch.device("meta"): + model = cls( + video_cross_attention_dim=int(transformer.get("cross_attention_dim", 4096)), + audio_cross_attention_dim=int(transformer.get("audio_cross_attention_dim", 2048)), + pooler_hidden_dim=int(duration.get("pooler_hidden_dim", 256)), + num_queries=int(duration.get("num_queries", 1)), + num_pooler_heads=int(duration.get("num_pooler_heads", 4)), + mlp_hidden=int(duration.get("mlp_hidden", 256)), + ) + unexpected, missing = ltx25_duration_checkpoint_key_coverage(checkpoint_path, set(model.state_dict())) + if unexpected or missing: + raise ValueError( + "LTX-2.5 DurationHead checkpoint coverage mismatch: " + f"unexpected={sorted(unexpected)[:5]}, missing={sorted(missing)[:5]}" + ) + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as checkpoint: + state_dict = { + key.removeprefix("duration_head."): checkpoint.get_tensor(key) + for key in checkpoint.keys() + if key.startswith("duration_head.") + } + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=True, assign=True) + if missing_keys or unexpected_keys: + raise ValueError( + f"LTX-2.5 DurationHead load mismatch: missing={missing_keys}, unexpected={unexpected_keys}" + ) + return model.to(device=device, dtype=torch_dtype).eval() + + +def ltx25_duration_checkpoint_key_coverage( + checkpoint_path: str | Path, model_keys: set[str] +) -> tuple[set[str], set[str]]: + """Return unexplained source keys and missing isolated DurationHead keys.""" + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as checkpoint: + mapped = {key.removeprefix("duration_head.") for key in checkpoint.keys() if key.startswith("duration_head.")} + return mapped - model_keys, model_keys - mapped + + +def seconds_to_num_frames( + seconds: float, + *, + frame_rate: float, + min_seconds: float = 1.0, + max_seconds: float = 20.0, +) -> int: + """Clamp a duration then snap it upward to LTX's causal ``8k + 1`` frame grid.""" + if frame_rate <= 0: + raise ValueError("frame_rate must be positive") + min_frames = round(min_seconds * frame_rate) + max_frames = round(max_seconds * frame_rate) + raw_frames = min(max(round(seconds * frame_rate), min_frames), max_frames) + frames = ((raw_frames - 1) // 8) * 8 + 1 + if frames < min_frames: + frames = min(((min_frames - 1 + 7) // 8) * 8 + 1, max_frames) + return frames + + +__all__ = [ + "LTX25DurationHead", + "ltx25_duration_checkpoint_key_coverage", + "seconds_to_num_frames", +] diff --git a/telefuser/models/ltx25/embeddings.py b/telefuser/models/ltx25/embeddings.py new file mode 100644 index 0000000..8d81348 --- /dev/null +++ b/telefuser/models/ltx25/embeddings.py @@ -0,0 +1,318 @@ +"""LTX-2.5 Gemma feature extraction and audio/video embeddings connectors.""" + +from __future__ import annotations + +import math +from pathlib import Path +from typing import Any, NamedTuple + +import torch +from safetensors import safe_open + +from telefuser.core.config import AttentionConfig, AttnImplType + +from .checkpoint import inspect_checkpoint +from .gemma4 import LTX25GemmaAssets +from .transformer import ( + Attention, + FeedForward, + LTXRopeType, + generate_freq_grid_np, + generate_freq_grid_pytorch, + precompute_freqs_cis, + rms_norm, +) + +_TRANSFORMER_PREFIX = "model.diffusion_model." +_VIDEO_CONNECTOR_PREFIX = _TRANSFORMER_PREFIX + "video_embeddings_connector." +_AUDIO_CONNECTOR_PREFIX = _TRANSFORMER_PREFIX + "audio_embeddings_connector." +_TEXT_PROJECTION_PREFIX = "text_embedding_projection." + + +class LTX25EmbeddingsProcessorOutput(NamedTuple): + """Conditioning tensors consumed by the video and audio diffusion paths.""" + + video_encoding: torch.Tensor + audio_encoding: torch.Tensor + attention_mask: torch.Tensor + + +def _right_pad_order(additive_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + binary = (additive_mask[:, 0, 0, :] >= 0).to(torch.int32) + indices = torch.argsort(binary, dim=-1, descending=True, stable=True) + ordered = torch.gather(binary, 1, indices) + mask = (ordered.to(additive_mask.dtype) - 1) * torch.finfo(additive_mask.dtype).max + return indices, mask[:, None, None, :] + + +def _apply_order(features: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: + return torch.gather(features, 1, indices.unsqueeze(-1).expand_as(features)) + + +def _additive_mask(mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + return (mask.to(torch.int64) - 1).to(dtype).reshape(mask.shape[0], 1, 1, mask.shape[-1]) * torch.finfo(dtype).max + + +class LTX25FeatureExtractor(torch.nn.Module): + """LTX-2.5 per-token Gemma RMS feature extraction and dual projections.""" + + def __init__(self, hidden_size: int, num_hidden_layers: int, video_dim: int, audio_dim: int) -> None: + super().__init__() + self.embedding_dim = hidden_size + self.flat_dim = hidden_size * (num_hidden_layers + 1) + self.video_aggregate_embed = torch.nn.Linear(self.flat_dim, video_dim, bias=True) + self.audio_aggregate_embed = torch.nn.Linear(self.flat_dim, audio_dim, bias=True) + + def forward( + self, + hidden_states: tuple[torch.Tensor, ...], + attention_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + encoded = torch.stack(hidden_states, dim=-1) + variance = torch.mean(encoded.square(), dim=2, keepdim=True) + normalized = encoded * torch.rsqrt(variance + 1e-6) + normalized = normalized.reshape(encoded.shape[0], encoded.shape[1], -1).to(encoded.dtype) + normalized = torch.where(attention_mask.bool().unsqueeze(-1), normalized, torch.zeros_like(normalized)) + video = self.video_aggregate_embed( + normalized * math.sqrt(self.video_aggregate_embed.out_features / self.embedding_dim) + ) + audio = self.audio_aggregate_embed( + normalized * math.sqrt(self.audio_aggregate_embed.out_features / self.embedding_dim) + ) + return video, audio + + +class _LTX25ConnectorBlock(torch.nn.Module): + """Pre-norm RoPE attention and feed-forward block used by each connector.""" + + def __init__(self, dim: int, heads: int, dim_head: int, rope_type: LTXRopeType, gated_attention: bool) -> None: + super().__init__() + self.attn1 = Attention( + query_dim=dim, + heads=heads, + dim_head=dim_head, + rope_type=rope_type, + apply_gated_attention=gated_attention, + ) + # Connector feed-forward layers retain the checkpoint's biases. + self.ff = FeedForward(dim, dim_out=dim, bias=True) + + def forward( + self, + hidden_states: torch.Tensor, + additive_attention_mask: torch.Tensor, + positional_embeddings: torch.Tensor, + ) -> torch.Tensor: + hidden_states = hidden_states + self.attn1( + rms_norm(hidden_states), mask=additive_attention_mask, pe=positional_embeddings + ) + return hidden_states + self.ff(rms_norm(hidden_states)) + + +class LTX25EmbeddingsConnector(torch.nn.Module): + """LTX-2.5 learned-register, 1D text-conditioning connector.""" + + def __init__( + self, + *, + attention_head_dim: int, + num_attention_heads: int, + num_layers: int, + positional_embedding_max_pos: list[int], + rope_type: LTXRopeType, + double_precision_rope: bool, + apply_gated_attention: bool, + num_learnable_registers: int, + ) -> None: + super().__init__() + self.inner_dim = num_attention_heads * attention_head_dim + self.num_attention_heads = num_attention_heads + self.positional_embedding_max_pos = positional_embedding_max_pos + self.rope_type = rope_type + self.double_precision_rope = double_precision_rope + self.num_learnable_registers = num_learnable_registers + self.transformer_1d_blocks = torch.nn.ModuleList( + [ + _LTX25ConnectorBlock( + self.inner_dim, + num_attention_heads, + attention_head_dim, + rope_type, + apply_gated_attention, + ) + for _ in range(num_layers) + ] + ) + self.learnable_registers = torch.nn.Parameter(torch.empty(num_learnable_registers, self.inner_dim)) + + def forward( + self, + hidden_states: torch.Tensor, + additive_attention_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, sequence_length, _ = hidden_states.shape + if sequence_length % self.num_learnable_registers: + raise ValueError( + "LTX-2.5 connector token length must be divisible by its learnable-register count: " + f"{sequence_length} % {self.num_learnable_registers}" + ) + registers = self.learnable_registers.repeat(sequence_length // self.num_learnable_registers, 1) + registers = registers.to(hidden_states).unsqueeze(0).expand(batch_size, -1, -1) + valid = (additive_attention_mask[:, 0, 0, :].unsqueeze(-1) >= 0).to(hidden_states.dtype) + hidden_states = valid * hidden_states + (1 - valid) * registers + connector_mask = torch.zeros_like(additive_attention_mask) + indices = torch.arange(sequence_length, dtype=torch.float32, device=hidden_states.device) + indices = indices[None, None, :].expand(batch_size, -1, -1) + generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch + positional_embeddings = precompute_freqs_cis( + indices, + dim=self.inner_dim, + out_dtype=hidden_states.dtype, + theta=10000.0, + max_pos=self.positional_embedding_max_pos, + num_attention_heads=self.num_attention_heads, + rope_type=self.rope_type, + freq_grid_generator=generator, + ) + for block in self.transformer_1d_blocks: + hidden_states = block(hidden_states, connector_mask, positional_embeddings) + return rms_norm(hidden_states), connector_mask + + +class LTX25EmbeddingsProcessor(torch.nn.Module): + """Feature extraction and dual connectors loaded from LTX-2.5 split checkpoints.""" + + def __init__(self, transformer_config: dict[str, Any], gemma_config: dict[str, Any]) -> None: + super().__init__() + text_config = gemma_config.get("text_config") + if not isinstance(text_config, dict): + raise ValueError("LTX-2.5 Gemma config is missing object text_config") + hidden_size = text_config.get("hidden_size") + num_hidden_layers = text_config.get("num_hidden_layers") + if not isinstance(hidden_size, int) or not isinstance(num_hidden_layers, int): + raise ValueError("LTX-2.5 Gemma text_config is missing hidden_size or num_hidden_layers") + rope_type = LTXRopeType(transformer_config["rope_type"]) + double_precision_rope = transformer_config.get("frequencies_precision") == "float64" + common = { + "num_layers": transformer_config["connector_num_layers"], + "positional_embedding_max_pos": transformer_config["connector_positional_embedding_max_pos"], + "rope_type": rope_type, + "double_precision_rope": double_precision_rope, + "apply_gated_attention": transformer_config["connector_apply_gated_attention"], + "num_learnable_registers": transformer_config["connector_num_learnable_registers"], + } + video_dim = transformer_config["num_attention_heads"] * transformer_config["attention_head_dim"] + audio_dim = transformer_config["audio_num_attention_heads"] * transformer_config["audio_attention_head_dim"] + self.feature_extractor = LTX25FeatureExtractor(hidden_size, num_hidden_layers, video_dim, audio_dim) + self.video_connector = LTX25EmbeddingsConnector( + attention_head_dim=transformer_config["connector_attention_head_dim"], + num_attention_heads=transformer_config["connector_num_attention_heads"], + **common, + ) + self.audio_connector = LTX25EmbeddingsConnector( + attention_head_dim=transformer_config["audio_connector_attention_head_dim"], + num_attention_heads=transformer_config["audio_connector_num_attention_heads"], + **common, + ) + + def forward( + self, + hidden_states: tuple[torch.Tensor, ...], + attention_mask: torch.Tensor, + ) -> LTX25EmbeddingsProcessorOutput: + # Upstream connector blocks use PyTorch's native SDPA priority path. + # Keep that exact baseline locally while leaving the denoiser's public + # attention dispatch configurable by its runtime configuration. + native_attention_config = AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA) + connector_attentions = [ + *(block.attn1 for block in self.video_connector.transformer_1d_blocks), + *(block.attn1 for block in self.audio_connector.transformer_1d_blocks), + ] + original_attention_configs = [attention.attention_config for attention in connector_attentions] + for attention in connector_attentions: + attention.attention_config = native_attention_config + if any(attention.attention_config is not native_attention_config for attention in connector_attentions): + raise RuntimeError("LTX-2.5 connector attention configuration was not applied") + try: + video_features, audio_features = self.feature_extractor(hidden_states, attention_mask) + additive_mask = _additive_mask(attention_mask, video_features.dtype) + order, connector_mask = _right_pad_order(additive_mask) + video_encoding, output_mask = self.video_connector(_apply_order(video_features, order), connector_mask) + audio_encoding, _ = self.audio_connector(_apply_order(audio_features, order), connector_mask) + binary_mask = (output_mask[:, 0, 0, :] >= 0).to(torch.int64) + return LTX25EmbeddingsProcessorOutput( + video_encoding * binary_mask.unsqueeze(-1), audio_encoding, binary_mask + ) + finally: + for attention, original_attention_config in zip( + connector_attentions, original_attention_configs, strict=True + ): + attention.attention_config = original_attention_config + + @classmethod + def from_checkpoints( + cls, + transformer_path: str | Path, + text_encoder_path: str | Path, + *, + device: torch.device | str = "cpu", + torch_dtype: torch.dtype = torch.bfloat16, + ) -> "LTX25EmbeddingsProcessor": + """Construct and strictly load connector weights from the official split files.""" + transformer = inspect_checkpoint(transformer_path) + assets = LTX25GemmaAssets.load(text_encoder_path) + transformer_config = transformer.config.get("transformer") + if not isinstance(transformer_config, dict): + raise ValueError("LTX-2.5 transformer checkpoint is missing object transformer config") + with torch.device("meta"): + model = cls(transformer_config, assets.config) + unexpected, missing = embeddings_checkpoint_key_coverage( + transformer.path, assets.checkpoint_path, set(model.state_dict()) + ) + if unexpected or missing: + raise ValueError( + "LTX-2.5 embeddings checkpoint coverage mismatch: " + f"unexpected={sorted(unexpected)[:5]}, missing={sorted(missing)[:5]}" + ) + state_dict = _load_embedding_state_dict(transformer.path, assets.checkpoint_path) + missing, unexpected = model.load_state_dict(state_dict, strict=True, assign=True) + if missing or unexpected: + raise ValueError(f"LTX-2.5 embeddings load mismatch: missing={missing[:5]}, unexpected={unexpected[:5]}") + return model.to(device=device, dtype=torch_dtype).eval() + + +def _embedding_target_key(path: Path, key: str) -> str | None: + if path.name.startswith("ltx-2.5-22b"): + if key.startswith(_VIDEO_CONNECTOR_PREFIX): + return "video_connector." + key.removeprefix(_VIDEO_CONNECTOR_PREFIX) + if key.startswith(_AUDIO_CONNECTOR_PREFIX): + return "audio_connector." + key.removeprefix(_AUDIO_CONNECTOR_PREFIX) + if key.startswith(_TEXT_PROJECTION_PREFIX): + return "feature_extractor." + key.removeprefix(_TEXT_PROJECTION_PREFIX) + return None + + +def _load_embedding_state_dict(transformer_path: Path, text_encoder_path: Path) -> dict[str, torch.Tensor]: + state_dict: dict[str, torch.Tensor] = {} + for path in (transformer_path, text_encoder_path): + with safe_open(str(path), framework="pt", device="cpu") as checkpoint: + for key in checkpoint.keys(): + target = _embedding_target_key(path, key) + if target is not None: + state_dict[target] = checkpoint.get_tensor(key) + return state_dict + + +def embeddings_checkpoint_key_coverage( + transformer_path: str | Path, + text_encoder_path: str | Path, + model_keys: set[str], +) -> tuple[set[str], set[str]]: + """Return unexplained source keys and missing model keys without materializing weights.""" + mapped: set[str] = set() + for path in (Path(transformer_path), Path(text_encoder_path)): + with safe_open(str(path), framework="pt", device="cpu") as checkpoint: + mapped.update( + target for key in checkpoint.keys() if (target := _embedding_target_key(path, key)) is not None + ) + return mapped - model_keys, model_keys - mapped diff --git a/telefuser/models/ltx25/gemma4.py b/telefuser/models/ltx25/gemma4.py new file mode 100644 index 0000000..76c0210 --- /dev/null +++ b/telefuser/models/ltx25/gemma4.py @@ -0,0 +1,272 @@ +"""Gemma4 Unified text-encoder loading for the LTX-2.5 packed checkpoint.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from safetensors import safe_open +from safetensors.torch import load_file +from tokenizers import Tokenizer +from transformers import AutoModelForImageTextToText, PreTrainedTokenizerFast +from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS +from transformers.models.auto.configuration_auto import CONFIG_MAPPING + +PACKED_GEMMA_CONFIG_KEY = "gemma_config" +PACKED_TOKENIZER_KEY = "tokenizer_json" +PACKED_ASSET_PREFIX = "hf_asset__" +GEMMA_MAX_LENGTH = 1024 +_CASTABLE_FLOAT_DTYPES = frozenset({torch.float16, torch.bfloat16, torch.float32, torch.float64}) + + +def _initialize_gemma4_unified_buffers(model: torch.nn.Module) -> None: + """Recompute non-persistent Gemma4 buffers after meta construction. + + The packed checkpoint stores model parameters only. Upstream rebuilds the + per-layer rotary frequencies and embedding scale after construction, rather + than relying on the meta-initialized values inherited from Transformers. + """ + language_model = model.model.language_model + config = model.config.text_config + rotary_embedding = language_model.rotary_emb + for layer_type in dict.fromkeys(config.layer_types): + rope_params = config.rope_parameters[layer_type] + if rope_params is None: + continue + rope_type = rope_params["rope_type"] + if rope_type == "default": + inv_freq, attention_scaling = rotary_embedding.compute_default_rope_parameters( + config, layer_type=layer_type + ) + else: + kwargs = {"layer_type": layer_type} + if layer_type == "full_attention" and rope_type == "proportional": + kwargs["head_dim_key"] = "global_head_dim" + inv_freq, attention_scaling = ROPE_INIT_FUNCTIONS[rope_type](config, **kwargs) + rotary_embedding.register_buffer(f"{layer_type}_inv_freq", inv_freq, persistent=False) + rotary_embedding.register_buffer(f"{layer_type}_original_inv_freq", inv_freq.clone(), persistent=False) + setattr(rotary_embedding, f"{layer_type}_attention_scaling", attention_scaling) + + language_model.embed_tokens.register_buffer( + "embed_scale", torch.tensor(config.hidden_size**0.5, device="cpu"), persistent=False + ) + if hasattr(language_model, "embed_tokens_per_layer"): + language_model.embed_tokens_per_layer.register_buffer( + "embed_scale", torch.tensor(config.hidden_size_per_layer_input**0.5, device="cpu"), persistent=False + ) + + +def _tensor_to_bytes(tensor: torch.Tensor) -> bytes: + """Decode a packed byte tensor without depending on its signedness.""" + array = tensor.detach().cpu().numpy() + return array.tobytes() if array.dtype == np.uint8 else array.astype(np.uint8).tobytes() + + +def _cast_checkpoint_tensor(tensor: torch.Tensor, torch_dtype: torch.dtype) -> torch.Tensor: + """Match the upstream builder's pre-assignment floating-point cast policy.""" + if tensor.dtype not in _CASTABLE_FLOAT_DTYPES: + return tensor + # Scalar float32 values are checkpoint scales, not model weights. Upstream + # keeps them in float32 to avoid losing scale precision. + if tensor.ndim == 0 and tensor.dtype == torch.float32: + return tensor + return tensor.to(dtype=torch_dtype) + + +@dataclass(frozen=True, slots=True) +class LTX25GemmaAssets: + """Hugging Face assets embedded in an LTX-2.5 Gemma4 checkpoint.""" + + checkpoint_path: Path + config: dict[str, Any] + tokenizer_json: bytes + sidecars: dict[str, bytes] + + @classmethod + def load(cls, checkpoint_path: str | Path) -> "LTX25GemmaAssets": + """Read embedded assets without materializing model weights.""" + path = Path(checkpoint_path).expanduser().resolve() + with safe_open(path, framework="pt", device="cpu") as checkpoint: + metadata = checkpoint.metadata() or {} + serialized_config = metadata.get(PACKED_GEMMA_CONFIG_KEY) + if serialized_config is None: + raise ValueError(f"LTX-2.5 Gemma checkpoint is missing {PACKED_GEMMA_CONFIG_KEY!r}: {path}") + config = json.loads(serialized_config) + keys = set(checkpoint.keys()) + if PACKED_TOKENIZER_KEY not in keys: + raise ValueError(f"LTX-2.5 Gemma checkpoint is missing {PACKED_TOKENIZER_KEY!r}: {path}") + tokenizer_json = _tensor_to_bytes(checkpoint.get_tensor(PACKED_TOKENIZER_KEY)) + sidecars = { + key.removeprefix(PACKED_ASSET_PREFIX): _tensor_to_bytes(checkpoint.get_tensor(key)) + for key in keys + if key.startswith(PACKED_ASSET_PREFIX) + } + if config.get("model_type") != "gemma4_unified": + raise ValueError( + f"LTX-2.5 requires model_type='gemma4_unified', got {config.get('model_type')!r} from {path}" + ) + if "tokenizer_config.json" not in sidecars: + raise ValueError(f"LTX-2.5 Gemma checkpoint is missing tokenizer_config.json: {path}") + return cls(path, config, tokenizer_json, sidecars) + + def build_config(self) -> Any: + """Build the registered Transformers configuration from packed JSON.""" + model_type = self.config.get("model_type") + try: + return CONFIG_MAPPING[model_type].from_dict(self.config) + except KeyError as exc: + raise ValueError(f"Unsupported packed Gemma model_type={model_type!r}") from exc + + def build_tokenizer(self) -> PreTrainedTokenizerFast: + """Build the LTX Gemma tokenizer with upstream packed sidecar settings.""" + tokenizer_config = json.loads(self.sidecars["tokenizer_config.json"]) + excluded = { + "tokenizer_class", + "auto_map", + "model_max_length", + "backend", + "is_local", + "local_files_only", + "processor_class", + "added_tokens_decoder", + } + kwargs = {key: value for key, value in tokenizer_config.items() if key not in excluded} + template = self.sidecars.get("chat_template.jinja") + if template is not None: + kwargs.setdefault("chat_template", template.decode()) + return PreTrainedTokenizerFast( + tokenizer_object=Tokenizer.from_buffer(self.tokenizer_json), + model_max_length=GEMMA_MAX_LENGTH, + **kwargs, + ) + + +class LTX25GemmaTokenizer: + """LTX's 1024-token, left-padded Gemma encoding adapter.""" + + def __init__(self, tokenizer: PreTrainedTokenizerFast) -> None: + self.tokenizer = tokenizer + self.tokenizer.model_max_length = GEMMA_MAX_LENGTH + self.tokenizer.padding_side = "left" + if self.tokenizer.pad_token is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + + def encode(self, prompts: list[str], device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + """Tokenize prompts using the upstream BOS insertion and left-padding contract.""" + bos_token_id = self.tokenizer.bos_token_id + if bos_token_id is None: + raise ValueError("LTX-2.5 Gemma tokenizer is missing bos_token_id") + token_ids: list[list[int]] = [] + for prompt in prompts: + encoded = self.tokenizer( + prompt.strip(), + padding=False, + truncation=True, + max_length=GEMMA_MAX_LENGTH, + return_tensors="pt", + ) + values = encoded.input_ids[0].tolist() + if not values or values[0] != bos_token_id: + values = [bos_token_id, *values][:GEMMA_MAX_LENGTH] + token_ids.append(values) + padded = self.tokenizer.pad( + {"input_ids": token_ids}, + padding="max_length", + max_length=GEMMA_MAX_LENGTH, + return_tensors="pt", + return_attention_mask=True, + ) + return padded.input_ids.to(device), padded.attention_mask.to(device) + + +def gemma4_checkpoint_key_to_model_key(key: str) -> str | None: + """Map packed Comfy-flat Gemma4 keys into ``LTX25Gemma4TextEncoder`` keys.""" + if key == PACKED_TOKENIZER_KEY or key.startswith(PACKED_ASSET_PREFIX): + return None + if key.startswith("model.layers."): + return "model.model.language_model." + key.removeprefix("model.") + if key.startswith("model.embed_tokens.") or key.startswith("model.norm."): + return "model.model.language_model." + key.removeprefix("model.") + if key.startswith("vision_model."): + return "model.model.embed_vision." + key.removeprefix("vision_model.") + if key.startswith("multi_modal_projector.embedding_projection."): + suffix = key.removeprefix("multi_modal_projector.embedding_projection.") + return "model.model.embed_vision.multimodal_embedder.embedding_projection." + suffix + if key.startswith("audio_projector."): + return "model.model.embed_audio." + key.removeprefix("audio_projector.") + return None + + +def gemma4_checkpoint_key_coverage(checkpoint_path: str | Path, model_keys: set[str]) -> tuple[set[str], set[str]]: + """Return unexpected mapped checkpoint keys and missing non-tied model keys.""" + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as checkpoint: + source_keys = checkpoint.keys() + mapped = {target for key in source_keys if (target := gemma4_checkpoint_key_to_model_key(key)) is not None} + expected = set(model_keys) + expected.discard("model.lm_head.weight") + return mapped - expected, expected - mapped + + +class LTX25Gemma4TextEncoder(torch.nn.Module): + """Gemma4 Unified wrapper that exposes LTX's raw hidden-state encoding path.""" + + def __init__(self, model: torch.nn.Module, tokenizer: LTX25GemmaTokenizer) -> None: + super().__init__() + self.model = model + self.tokenizer = tokenizer + + @classmethod + def from_checkpoint( + cls, + checkpoint_path: str | Path, + *, + device: torch.device | str = "cpu", + torch_dtype: torch.dtype = torch.bfloat16, + strict: bool = True, + ) -> "LTX25Gemma4TextEncoder": + """Construct and assign the packed Gemma4 checkpoint with explicit key coverage.""" + assets = LTX25GemmaAssets.load(checkpoint_path) + config = assets.build_config() + # Match the upstream Gemma4 configurator: parameters and transient + # buffers originate on meta, then the packed checkpoint assigns weights. + with torch.device("meta"): + model = AutoModelForImageTextToText.from_config(config) + _initialize_gemma4_unified_buffers(model) + wrapper = cls(model, LTX25GemmaTokenizer(assets.build_tokenizer())) + unexpected, missing = gemma4_checkpoint_key_coverage(checkpoint_path, set(wrapper.state_dict())) + if unexpected or missing: + raise ValueError( + "LTX-2.5 Gemma checkpoint coverage mismatch: " + f"unexpected={sorted(unexpected)[:5]}, missing={sorted(missing)[:5]}" + ) + source = load_file(str(assets.checkpoint_path), device="cpu") + state_dict = { + target: _cast_checkpoint_tensor(tensor, torch_dtype) + for key, tensor in source.items() + if (target := gemma4_checkpoint_key_to_model_key(key)) is not None + } + embed_key = "model.model.language_model.embed_tokens.weight" + state_dict["model.lm_head.weight"] = state_dict[embed_key] + missing_keys, unexpected_keys = wrapper.load_state_dict(state_dict, strict=strict, assign=True) + if missing_keys or unexpected_keys: + raise ValueError( + f"LTX-2.5 Gemma load mismatch: missing={missing_keys[:5]}, unexpected={unexpected_keys[:5]}" + ) + # The upstream builder casts checkpoint weights before assignment and + # then transfers the module without a dtype conversion. In particular, + # Gemma's non-persistent rotary buffers remain float32. + return wrapper.to(device=device).eval() + + @torch.inference_mode() + def encode(self, prompts: list[str]) -> tuple[tuple[torch.Tensor, ...], torch.Tensor, torch.Tensor]: + """Return the raw Gemma hidden states, token IDs, and binary attention mask.""" + if not prompts: + raise ValueError("LTX-2.5 Gemma encode requires at least one prompt") + device = next(self.model.parameters()).device + input_ids, attention_mask = self.tokenizer.encode(prompts, device) + outputs = self.model.model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True) + return tuple(outputs.hidden_states), input_ids, attention_mask diff --git a/telefuser/models/ltx25/memory_efficient_decode.py b/telefuser/models/ltx25/memory_efficient_decode.py new file mode 100644 index 0000000..d1e9e55 --- /dev/null +++ b/telefuser/models/ltx25/memory_efficient_decode.py @@ -0,0 +1,638 @@ +"""Memory-efficient VAE decoder operations. +Reduces peak VRAM usage during video decoding through in-place operations +and workspace buffer reuse. The main optimizations are: +1. **Workspace buffers** -- Pre-allocated tensors with temporal padding replace + dynamic padding (``F.pad`` / ``concatenate``) in ``CausalConv3d``. A + workspace of shape ``[B, C, T+2, H, W]`` holds the data in positions + ``[1:-1]`` with replicate padding at ``[0]`` and ``[-1]``. +2. **In-place temporal-chunked Conv3d** *(non-causal only)* -- The convolution + output is written back into the workspace buffer, avoiding a separate + output allocation. Temporal chunking with boundary save/restore ensures + correct reads despite in-place writes. +3. **In-place normalization and affine transforms** -- PixelNorm, scale/shift, + and SiLU are applied in-place on workspace views. +4. **Free-before-conv** -- For ``DepthToSpaceUpsample`` blocks the input + tensor is freed before the convolution runs so that peak VRAM never holds + input *and* output simultaneously. +Both causal and non-causal modes are supported. Non-causal mode benefits +from all four optimizations. Causal mode benefits from optimizations 1, 3, +and 4; in-place conv (2) is skipped because the asymmetric causal padding +layout prevents clean in-place overwrites. +Use :func:`convert_decoder_weights_to_channels_last_3d` before +:func:`enable_memory_efficient_decode` to match the official decoder path. +""" + +from __future__ import annotations + +import math + +import torch +from einops import rearrange +from torch import nn +from torch.nn import functional as F + +from .conv_video_vae import ( + CausalConv3d, + DepthToSpaceUpsample, + PixelNorm, + ResnetBlock3D, + UNetMidBlock3D, + VideoDecoder, + unpatchify, +) + +# --------------------------------------------------------------------------- +# Low-level helpers +# --------------------------------------------------------------------------- + + +def _memory_format_of(t: torch.Tensor, prefer_channels_last_3d: bool = False) -> torch.memory_format: + """Pick the memory format for a workspace allocation. + When ``prefer_channels_last_3d`` is True and ``t`` is 5D, return + ``channels_last_3d`` regardless of ``t``'s current strides -- the + workspace's ``.copy_(t)`` will transcribe the data into the new layout. + This is needed because intermediate tensors inside the decoder (after + ``rearrange`` + slice + residual add in ``_upsample_forward_efficient``) + are not NHWC-contiguous, so an auto-detect helper would silently fall + back to NCHW for every workspace after the first upsample. + Otherwise fall back to inspecting ``t``: ``channels_last_3d`` if ``t`` + already uses it, else contiguous. + """ + if prefer_channels_last_3d and t.dim() == 5: + return torch.channels_last_3d + if t.dim() == 5 and t.is_contiguous(memory_format=torch.channels_last_3d): + return torch.channels_last_3d + return torch.contiguous_format + + +def _find_temporal_split_size(num_frames: int) -> int: + """Find chunk size for in-place temporal convolution. + The chunk size ensures the last chunk has at least 3 frames + (the temporal kernel size), avoiding degenerate chunks. + """ + for s in range(16, 2, -1): + remainder = num_frames % s + if remainder == 0 or remainder >= 3: + return s + + raise ValueError( + f"Unable to find a valid temporal split size for num_frames={num_frames}. " + "Expected a split size between 3 and 16 such that the final chunk is " + "either exact or has at least 3 frames." + ) + + +def _pad_workspace_temporal(workspace: torch.Tensor) -> None: + """Apply non-causal replicate padding to temporal boundaries. + Sets ``workspace[:, :, 0]`` to a copy of ``workspace[:, :, 1]`` and + ``workspace[:, :, -1]`` to a copy of ``workspace[:, :, -2]``. + """ + workspace[:, :, 0, :, :].copy_(workspace[:, :, 1, :, :]) + workspace[:, :, -1, :, :].copy_(workspace[:, :, -2, :, :]) + + +# --------------------------------------------------------------------------- +# In-place Conv3d (non-causal only) +# --------------------------------------------------------------------------- + + +def inplace_conv3d_temporal_chunked(workspace: torch.Tensor, conv: nn.Conv3d) -> None: + """Run a 3x3x3 Conv3d in-place on a temporally-padded workspace. + The workspace has shape ``[B, C, T+2, H, W]`` where positions ``[1:-1]`` + hold the real data and positions ``[0]`` and ``[-1]`` are padding slots. + The convolution must have ``kernel_size=(3,3,3)``, ``stride=(1,1,1)``, + ``padding=(0,1,1)`` -- no temporal padding, symmetric spatial padding. + The output (T frames) overwrites positions ``[1:-1]``. Temporal chunking + with boundary save/restore ensures each chunk reads unmodified input even + though earlier chunks already wrote to the same buffer. + Only valid for **non-causal** mode (symmetric replicate padding). + Args: + workspace: Tensor ``[B, max(C_in, C_out), T+2, H, W]``. + Modified in-place; after the call ``workspace[:, :C_out, 1:-1]`` + holds the convolution result. + conv: ``nn.Conv3d`` with the constraints above. + """ + if conv.kernel_size != (3, 3, 3): + raise ValueError(f"Expected kernel_size=(3,3,3), got {conv.kernel_size}") + if conv.stride != (1, 1, 1): + raise ValueError(f"Expected stride=(1,1,1), got {conv.stride}") + if conv.padding != (0, 1, 1): + raise ValueError(f"Expected padding=(0,1,1), got {conv.padding}") + + _pad_workspace_temporal(workspace) + + total_frames = workspace.shape[2] + out_channels = conv.out_channels + in_channels = conv.in_channels + + if total_frames > 16: + split_size = _find_temporal_split_size(total_frames) + num_splits = (total_frames + split_size - 1) // split_size + else: + split_size = total_frames - 1 + num_splits = 1 + + # 1-frame buffers for saving / restoring boundary frames across chunks. + x_buf = torch.empty( + workspace.shape[0], + workspace.shape[1], + 1, + workspace.shape[3], + workspace.shape[4], + device=workspace.device, + dtype=workspace.dtype, + memory_format=_memory_format_of(workspace), + ) + o_buf = torch.empty_like(x_buf) + + # Helper: extract a chunk and make it contiguous. Workspace views can + # inherit strides > 2^31 from the full buffer, which makes Conv3d's + # reflect-padding path (F.pad) crash with "input tensor must fit into + # 32-bit index math". A small .clone() per chunk avoids this. + needs_clone = workspace.untyped_storage().nbytes() > (2**31 - 1) * workspace.element_size() + + def _chunk(t_start: int, t_end: int) -> torch.Tensor: + s = workspace[:, :in_channels, t_start:t_end] + return s.clone() if needs_clone else s + + # --- First chunk --- + if num_splits > 1: + # Save the boundary now so the loop below can restore it. Skipped + # when there is only one chunk: the loop never runs, and the save + # would be a wasted full HW slice copy. + x_buf[:, :, 0] = workspace[:, :, split_size - 1].clone() + workspace[:, :out_channels, 1:split_size] = conv(_chunk(0, split_size + 1)) + + # --- Remaining chunks --- + for i in range(1, num_splits): + start = i * split_size + end = min((i + 1) * split_size, total_frames - 1) + + # Save the value at start-1 (now holds previous chunk's output). + o_buf[:, :, 0] = workspace[:, :, start - 1].clone() + # Restore the original input value needed by this chunk's conv. + workspace[:, :, start - 1] = x_buf[:, :, 0] + # Save the boundary for the *next* chunk before we overwrite it. + x_buf[:, :, 0] = workspace[:, :, end - 1].clone() + + workspace[:, :out_channels, start:end] = conv(_chunk(start - 1, end + 1)) + + # Put back the previous chunk's output at the boundary. + workspace[:, :, start - 1] = o_buf[:, :, 0] + + +# --------------------------------------------------------------------------- +# Causal conv helper (free-before-conv) +# --------------------------------------------------------------------------- + + +def _causal_pad(x: torch.Tensor, pad_size: int) -> torch.Tensor: + """Build a causal-padded buffer of shape ``[B, C, T+pad_size, H, W]``. + Copies ``x`` into ``padded[:, :, pad_size:]`` and replicates the first + real frame into the leading ``pad_size`` slots. The caller still owns + ``x`` after this returns. + """ + padded = torch.empty( + x.shape[0], + x.shape[1], + x.shape[2] + pad_size, + x.shape[3], + x.shape[4], + device=x.device, + dtype=x.dtype, + memory_format=_memory_format_of(x), + ) + padded[:, :, pad_size:].copy_(x) + for i in range(pad_size): + padded[:, :, i] = padded[:, :, pad_size] + return padded + + +def _causal_pad_free_and_conv(x: torch.Tensor, causal_conv: CausalConv3d) -> torch.Tensor: + """Causal-pad *x*, free it, then run the raw ``nn.Conv3d``. + This avoids the peak where both the original and padded tensors are + live simultaneously (as happens inside ``CausalConv3d.forward``). + Args: + x: Input ``[B, C_in, T, H, W]``. **Deleted** inside this function; + the caller must not use it afterwards. + Returns: + Convolution output ``[B, C_out, T, H, W]``. + """ + padded = _causal_pad(x, causal_conv.time_kernel_size - 1) + del x + result = causal_conv.conv(padded) + del padded + return result + + +# --------------------------------------------------------------------------- +# In-place normalization +# --------------------------------------------------------------------------- + + +def _pixel_norm_inplace(x: torch.Tensor, eps: float = 1e-8) -> None: + """In-place RMS (pixel) normalization along the channel dimension.""" + rms = torch.sqrt(torch.mean(x**2, dim=1, keepdim=True) + eps) + x.div_(rms) + + +def _norm_inplace(norm: nn.Module, x: torch.Tensor) -> None: + """Apply *norm* in-place, using an optimised path for ``PixelNorm``.""" + if isinstance(norm, PixelNorm): + _pixel_norm_inplace(x, eps=norm.eps) + else: + # GroupNorm or other -- fall back to allocating a temporary. + result = norm(x) + x.copy_(result) + del result + + +# --------------------------------------------------------------------------- +# Per-block efficient forwards +# --------------------------------------------------------------------------- + + +def _resnet_block_forward_inplace( + resnet: ResnetBlock3D, + workspace: torch.Tensor, + causal: bool, + timestep: torch.Tensor | None, + generator: torch.Generator | None, +) -> None: + """Run a ``ResnetBlock3D`` in-place on a workspace buffer. + The workspace has shape ``[B, C, T+2, H, W]`` with real data in + ``[1:-1]``. After this call ``workspace[:, :, 1:-1]`` holds the + residual-branch output ``F(x)`` (without the skip connection -- + the caller adds it back to the hidden state). + Only valid when ``in_channels == out_channels`` (true for all + ``ResnetBlock3D`` instances inside a ``UNetMidBlock3D``). + """ + if resnet.in_channels != resnet.out_channels: + raise ValueError( + "In-place resnet forward requires in_channels == out_channels, " + f"got {resnet.in_channels} != {resnet.out_channels}" + ) + + interior = workspace[:, :, 1:-1] + + # --- norm1 + [ada scaling] + SiLU + conv1 --- + _norm_inplace(resnet.norm1, interior) + + if resnet.timestep_conditioning and timestep is not None: + ada = resnet.scale_shift_table[None, ..., None, None, None].to( + device=interior.device, dtype=interior.dtype + ) + timestep.reshape( + interior.shape[0], + 4, + -1, + timestep.shape[-3], + timestep.shape[-2], + timestep.shape[-1], + ) + shift1, scale1, shift2, scale2 = ada.unbind(dim=1) + interior.mul_(1 + scale1).add_(shift1) + + F.silu(interior, inplace=True) + + if causal: + result = resnet.conv1(interior, causal=True) + interior.copy_(result) + del result + else: + inplace_conv3d_temporal_chunked(workspace, resnet.conv1.conv) + + if resnet.inject_noise: + spatial_shape = interior.shape[-2:] + scale = resnet.per_channel_scale1.to(device=interior.device, dtype=interior.dtype) + noise = torch.randn(spatial_shape, device=interior.device, dtype=interior.dtype, generator=generator) + interior.add_((noise * scale)[None, :, None, ...]) + + # --- norm2 + [ada scaling] + SiLU + conv2 --- + _norm_inplace(resnet.norm2, interior) + + if resnet.timestep_conditioning and timestep is not None: + interior.mul_(1 + scale2).add_(shift2) # type: ignore[possibly-undefined] + + F.silu(interior, inplace=True) + # dropout is always 0.0 during inference -- skip. + + if causal: + result = resnet.conv2(interior, causal=True) + interior.copy_(result) + del result + else: + inplace_conv3d_temporal_chunked(workspace, resnet.conv2.conv) + + if resnet.inject_noise: + spatial_shape = interior.shape[-2:] + scale = resnet.per_channel_scale2.to(device=interior.device, dtype=interior.dtype) + noise = torch.randn(spatial_shape, device=interior.device, dtype=interior.dtype, generator=generator) + interior.add_((noise * scale)[None, :, None, ...]) + + +def _midblock_forward_efficient( + block: UNetMidBlock3D, + hidden_states: torch.Tensor, + causal: bool, + timestep: torch.Tensor | None, + generator: torch.Generator | None, + prefer_channels_last_3d: bool = False, +) -> torch.Tensor: + """Memory-efficient ``UNetMidBlock3D`` forward. + Allocates a single workspace buffer that is reused across all + ``ResnetBlock3D`` iterations. For each block the workspace is + populated with the current hidden state, processed in-place, and + the result is added back (residual connection). + """ + timestep_embed = None + if block.timestep_conditioning: + if timestep is None: + raise ValueError("'timestep' required when timestep_conditioning=True") + batch_size = hidden_states.shape[0] + timestep_embed = block.time_embedder( + timestep=timestep.flatten(), + hidden_dtype=hidden_states.dtype, + ) + timestep_embed = timestep_embed.view(batch_size, timestep_embed.shape[-1], 1, 1, 1) + + workspace = torch.empty( + hidden_states.shape[0], + hidden_states.shape[1], + hidden_states.shape[2] + 2, + hidden_states.shape[3], + hidden_states.shape[4], + device=hidden_states.device, + dtype=hidden_states.dtype, + memory_format=_memory_format_of(hidden_states, prefer_channels_last_3d), + ) + + for resnet in block.res_blocks: + workspace[:, :, 1:-1].copy_(hidden_states) + _resnet_block_forward_inplace(resnet, workspace, causal, timestep_embed, generator) + hidden_states.add_(workspace[:, :, 1:-1]) + + del workspace + return hidden_states + + +def _upsample_forward_efficient( + block: DepthToSpaceUpsample, + x: torch.Tensor, + causal: bool, + prefer_channels_last_3d: bool = False, +) -> torch.Tensor: + """Memory-efficient ``DepthToSpaceUpsample`` forward. + For non-causal mode the input is copied into a workspace and the + convolution runs in-place. For causal mode the input is manually + padded and freed before the convolution runs. Both paths avoid + the peak where input *and* output coexist. + """ + if block.residual: + x_in = rearrange( + x, + "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)", + p1=block.stride[0], + p2=block.stride[1], + p3=block.stride[2], + ) + num_repeat = math.prod(block.stride) // block.out_channels_reduction_factor + x_in = x_in.repeat(1, num_repeat, 1, 1, 1) + if block.stride[0] == 2: + x_in = x_in[:, :, 1:, :, :] + + conv = block.conv.conv # underlying nn.Conv3d inside CausalConv3d + in_channels = x.shape[1] + out_channels = conv.out_channels + + if causal: + x = _causal_pad_free_and_conv(x, block.conv) + else: + mem_fmt = _memory_format_of(x, prefer_channels_last_3d) + workspace = torch.empty( + x.shape[0], + max(in_channels, out_channels), + x.shape[2] + 2, + x.shape[3], + x.shape[4], + device=x.device, + dtype=x.dtype, + memory_format=mem_fmt, + ) + workspace[:, :in_channels, 1:-1].copy_(x) + del x + inplace_conv3d_temporal_chunked(workspace, conv) + x = workspace[:, :out_channels, 1:-1].contiguous(memory_format=mem_fmt) + del workspace + + x = rearrange( + x, + "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)", + p1=block.stride[0], + p2=block.stride[1], + p3=block.stride[2], + ) + if block.stride[0] == 2: + x = x[:, :, 1:, :, :] + if block.residual: + x.add_(x_in) + del x_in + return x + + +# --------------------------------------------------------------------------- +# Final norm + conv_out +# --------------------------------------------------------------------------- + + +def _final_norm_and_conv_out( + decoder: VideoDecoder, + sample: torch.Tensor, + causal: bool, + scaled_timestep: torch.Tensor | None, + batch_size: int, + prefer_channels_last_3d: bool = False, +) -> torch.Tensor: + """Workspace-based final norm + [ada] + SiLU + conv_out + unpatchify.""" + conv_out_mod: CausalConv3d = decoder.conv_out # type: ignore[assignment] + conv_out = conv_out_mod.conv + feature_channels = sample.shape[1] + + mem_fmt = _memory_format_of(sample, prefer_channels_last_3d) + workspace = torch.empty( + sample.shape[0], + max(feature_channels, conv_out.out_channels), + sample.shape[2] + 2, + sample.shape[3], + sample.shape[4], + device=sample.device, + dtype=sample.dtype, + memory_format=mem_fmt, + ) + workspace[:, :feature_channels, 1:-1].copy_(sample) + del sample + + interior = workspace[:, :feature_channels, 1:-1] + _norm_inplace(decoder.conv_norm_out, interior) + + if decoder.timestep_conditioning: + embedded_timestep = decoder.last_time_embedder( + timestep=scaled_timestep.flatten(), + hidden_dtype=interior.dtype, + ) + embedded_timestep = embedded_timestep.view(batch_size, embedded_timestep.shape[-1], 1, 1, 1) + ada_values = decoder.last_scale_shift_table[None, ..., None, None, None].to( + device=interior.device, dtype=interior.dtype + ) + embedded_timestep.reshape( + batch_size, + 2, + -1, + embedded_timestep.shape[-3], + embedded_timestep.shape[-2], + embedded_timestep.shape[-1], + ) + shift, scale = ada_values.unbind(dim=1) + interior.mul_(1 + scale).add_(shift) + + F.silu(interior, inplace=True) + + if causal: + # Causal: build padded tensor directly from the interior view, + # then free the workspace before running the conv. + padded = _causal_pad(interior, conv_out_mod.time_kernel_size - 1) + del workspace, interior + result = conv_out(padded) + del padded + else: + inplace_conv3d_temporal_chunked(workspace, conv_out) + result = workspace[:, : conv_out.out_channels, 1:-1].contiguous(memory_format=mem_fmt) + del workspace, interior + + return unpatchify(result, patch_size_hw=decoder.patch_size, patch_size_t=1) + + +# --------------------------------------------------------------------------- +# Top-level efficient decoder forward +# --------------------------------------------------------------------------- + + +def _memory_efficient_forward( + decoder: VideoDecoder, + sample: torch.Tensor, + timestep: torch.Tensor | None = None, + generator: torch.Generator | None = None, +) -> torch.Tensor: + """Full memory-efficient ``VideoDecoder.forward`` replacement. + Orchestrates the entire decode through workspace-based operations: + ``UNetMidBlock3D`` and ``DepthToSpaceUpsample`` blocks use efficient + paths; standalone ``ResnetBlock3D`` blocks fall back to the standard + forward. The final norm + ada + SiLU + conv_out is also workspace-based. + All workspaces are allocated ``channels_last_3d`` so cuDNN's NHWC 3D + conv kernels run end-to-end. Weights must already be NHWC (via + :data:`CHANNELS_LAST_3D_WEIGHTS` chained at load); this path only formats the + input sample. + """ + causal = decoder.causal + batch_size = sample.shape[0] + sample = sample.to(next(decoder.parameters()).dtype) + + # --- Noise injection and de-normalisation (identical to standard path) --- + if decoder.timestep_conditioning: + noise = ( + torch.randn(sample.size(), generator=generator, dtype=sample.dtype, device=sample.device) + * decoder.decode_noise_scale + ) + sample = noise + (1.0 - decoder.decode_noise_scale) * sample + + sample = decoder.per_channel_statistics.un_normalize(sample) + + if timestep is None and decoder.timestep_conditioning: + timestep = torch.full((batch_size,), decoder.decode_timestep, device=sample.device, dtype=sample.dtype) + + # --- conv_in (latent tensor is small -- standard path is fine) --- + sample = decoder.conv_in(sample, causal=causal) + + upscale_dtype = next(iter(decoder.up_blocks.parameters())).dtype + sample = sample.to(upscale_dtype) + + scaled_timestep = None + if decoder.timestep_conditioning: + if timestep is None: + raise ValueError("'timestep' required when timestep_conditioning=True") + scaled_timestep = timestep * decoder.timestep_scale_multiplier.to(sample) + + # Workspaces are unconditionally NHWC: rearrange + slice + residual-add + # inside _upsample_forward_efficient produces NCHW-default output, so + # per-tensor inspection would silently fall back to NCHW for every + # workspace after the first upsample. + + # --- Up blocks (dispatch to efficient path per block type) --- + for up_block in decoder.up_blocks: + if isinstance(up_block, UNetMidBlock3D): + sample = _midblock_forward_efficient( + up_block, + sample, + causal=causal, + timestep=scaled_timestep if decoder.timestep_conditioning else None, + generator=generator, + prefer_channels_last_3d=True, + ) + elif isinstance(up_block, DepthToSpaceUpsample): + sample = _upsample_forward_efficient(up_block, sample, causal=causal, prefer_channels_last_3d=True) + elif isinstance(up_block, ResnetBlock3D): + sample = up_block(sample, causal=causal, generator=generator) + else: + sample = up_block(sample, causal=causal) + + return _final_norm_and_conv_out(decoder, sample, causal, scaled_timestep, batch_size, prefer_channels_last_3d=True) + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def convert_decoder_weights_to_channels_last_3d(decoder: VideoDecoder) -> VideoDecoder: + """Convert Conv3d decoder weights to the upstream channels-last layout.""" + for parameter in decoder.parameters(): + if parameter.dim() == 5: + parameter.data = parameter.data.to(memory_format=torch.channels_last_3d) + return decoder + + +def enable_memory_efficient_decode(decoder: nn.Module) -> nn.Module: + """Patch a ``VideoDecoder`` to use the memory-efficient forward path. + The mem-efficient path runs the decoder in ``channels_last_3d`` memory + format. **Weights** must already be NHWC — chain :data:`CHANNELS_LAST_3D_WEIGHTS` + into the builder's sd_ops at load so the registry stores that layout. This + patch only formats the **activation** and swaps in the workspace-based forward. + The original ``forward`` is saved as ``decoder._original_forward`` so + that it can be restored later with :func:`disable_memory_efficient_decode`. + """ + if not isinstance(decoder, VideoDecoder): + raise TypeError(f"Expected VideoDecoder, got {type(decoder).__name__}") + + if hasattr(decoder, "_original_forward"): + return decoder + + original_forward = decoder.forward + + def efficient_forward( + sample: torch.Tensor, + timestep: torch.Tensor | None = None, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + if sample.dim() == 5: + sample = sample.to(memory_format=torch.channels_last_3d) + return _memory_efficient_forward(decoder, sample, timestep, generator) + + decoder._original_forward = original_forward # type: ignore[attr-defined] + decoder.forward = efficient_forward # type: ignore[assignment] + return decoder + + +def disable_memory_efficient_decode(decoder: nn.Module) -> nn.Module: + """Restore the original ``forward`` method on a patched ``VideoDecoder``.""" + if hasattr(decoder, "_original_forward"): + decoder.forward = decoder._original_forward # type: ignore[attr-defined] + del decoder._original_forward # type: ignore[attr-defined] + return decoder diff --git a/telefuser/models/ltx25/sampler.py b/telefuser/models/ltx25/sampler.py new file mode 100644 index 0000000..b9451ef --- /dev/null +++ b/telefuser/models/ltx25/sampler.py @@ -0,0 +1,66 @@ +"""Exact LTX-2.5 distilled sampler constants and ancestral Euler step.""" + +from __future__ import annotations + +import torch + +from .checkpoint import parse_model_version + +LTX25_STAGE1_DISTILLED_SIGMAS = (1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0) +LTX25_STAGE2_DISTILLED_SIGMAS = (0.909375, 0.725, 0.421875, 0.0) +ANCESTRAL_NOISE_SEED_OFFSET = 10_000 + + +def distilled_sigmas(stage: int, *, device: torch.device | str | None = None) -> torch.Tensor: + """Create the upstream distilled sigma schedule as float32.""" + if stage == 1: + values = LTX25_STAGE1_DISTILLED_SIGMAS + elif stage == 2: + values = LTX25_STAGE2_DISTILLED_SIGMAS + else: + raise ValueError(f"stage must be 1 or 2, got {stage}") + return torch.tensor(values, dtype=torch.float32, device=device) + + +def uses_ancestral_stage1_sampler(model_version: str | None) -> bool: + """Return whether a checkpoint generation uses LTX-2.5 ancestral stage 1.""" + return parse_model_version(model_version) >= (2, 5) + + +class LTX25EulerAncestralStep: + """Upstream LTX-2.5 rectified-flow ancestral Euler update.""" + + def __init__(self, eta: float = 1.0, s_noise: float = 1.0) -> None: + self.eta = eta + self.s_noise = s_noise + + def step( + self, + sample: torch.Tensor, + denoised_sample: torch.Tensor, + sigmas: torch.Tensor, + step_index: int, + noise: torch.Tensor | None = None, + ) -> torch.Tensor: + """Advance one LTX rectified-flow ancestral Euler step.""" + sigma = sigmas[step_index].to(torch.float32) + sigma_next = sigmas[step_index + 1].to(torch.float32) + if bool(sigma_next == 0): + return denoised_sample.to(sample.dtype) + if self.eta > 0 and noise is None: + raise ValueError("LTX25EulerAncestralStep requires noise when eta > 0") + + sigma_down = sigma_next * (1.0 + (sigma_next / sigma - 1.0) * self.eta) + ratio = sigma_down / sigma + result = ratio * sample.float() + (1.0 - ratio) * denoised_sample.float() + if self.eta > 0: + alpha_next = 1.0 - sigma_next + alpha_down = 1.0 - sigma_down + renoise = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).clamp(min=0).sqrt() + result = (alpha_next / alpha_down) * result + noise.float() * self.s_noise * renoise + return result.to(sample.dtype) + + +def ancestral_noise_generator(seed: int, device: torch.device | str) -> torch.Generator: + """Create the independent upstream generator used for stage-1 ancestral noise.""" + return torch.Generator(device=device).manual_seed(seed + ANCESTRAL_NOISE_SEED_OFFSET) diff --git a/telefuser/models/ltx25/spatial_upsampler.py b/telefuser/models/ltx25/spatial_upsampler.py new file mode 100644 index 0000000..67bc3a6 --- /dev/null +++ b/telefuser/models/ltx25/spatial_upsampler.py @@ -0,0 +1,212 @@ +"""LTX-2.5 latent spatial upsampler and checkpoint loader.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +from safetensors import safe_open + +from .checkpoint import inspect_checkpoint + + +class LTX25PerChannelStatistics(torch.nn.Module): + """LTX video-latent normalization statistics stored in the VAE checkpoint.""" + + def __init__(self, channels: int) -> None: + super().__init__() + self.register_buffer("std-of-means", torch.ones(channels)) + self.register_buffer("mean-of-means", torch.zeros(channels)) + + def un_normalize(self, latent: torch.Tensor) -> torch.Tensor: + std = self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(latent) + mean = self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(latent) + return latent * std + mean + + def normalize(self, latent: torch.Tensor) -> torch.Tensor: + std = self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(latent) + mean = self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(latent) + return (latent - mean) / std + + +class LTX25UpsamplerResBlock(torch.nn.Module): + """The residual block used by the LTX-2.5 latent upsampler.""" + + def __init__(self, channels: int, dims: int) -> None: + super().__init__() + conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d + self.conv1 = conv(channels, channels, kernel_size=3, padding=1) + self.norm1 = torch.nn.GroupNorm(32, channels) + self.conv2 = conv(channels, channels, kernel_size=3, padding=1) + self.norm2 = torch.nn.GroupNorm(32, channels) + self.activation = torch.nn.SiLU() + + def forward(self, value: torch.Tensor) -> torch.Tensor: + residual = value + value = self.activation(self.norm1(self.conv1(value))) + value = self.norm2(self.conv2(value)) + return self.activation(value + residual) + + +class LTX25PixelShuffleND(torch.nn.Module): + """Channel-to-axis rearrangement matching the LTX-2.5 upsampler layout.""" + + def __init__(self, dims: int) -> None: + super().__init__() + if dims not in (1, 2, 3): + raise ValueError(f"Pixel shuffle dims must be 1, 2, or 3, got {dims}") + self.dims = dims + + def forward(self, value: torch.Tensor) -> torch.Tensor: + if self.dims == 2 and value.ndim == 4: + batch, channels, height, width = value.shape + if channels % 4: + raise ValueError(f"2D pixel shuffle requires channels divisible by 4, got {channels}") + value = value.reshape(batch, channels // 4, 2, 2, height, width) + return value.permute(0, 1, 4, 2, 5, 3).reshape(batch, channels // 4, height * 2, width * 2) + if value.ndim != 5: + raise ValueError(f"Pixel shuffle expects a 4D or 5D tensor, got {value.ndim}D") + batch, channels, frames, height, width = value.shape + if self.dims == 3: + if channels % 8: + raise ValueError(f"3D pixel shuffle requires channels divisible by 8, got {channels}") + value = value.reshape(batch, channels // 8, 2, 2, 2, frames, height, width) + return value.permute(0, 1, 5, 2, 6, 3, 7, 4).reshape( + batch, channels // 8, frames * 2, height * 2, width * 2 + ) + if self.dims == 2: + if channels % 4: + raise ValueError(f"2D pixel shuffle requires channels divisible by 4, got {channels}") + value = value.reshape(batch, channels // 4, 2, 2, frames, height, width) + return value.permute(0, 1, 4, 5, 2, 6, 3).reshape(batch, channels // 4, frames, height * 2, width * 2) + if channels % 2: + raise ValueError(f"1D pixel shuffle requires channels divisible by 2, got {channels}") + value = value.reshape(batch, channels // 2, 2, frames, height, width) + return value.permute(0, 1, 3, 2, 4, 5).reshape(batch, channels // 2, frames * 2, height, width) + + +@dataclass(frozen=True, slots=True) +class LTX25SpatialUpsamplerConfig: + """Architecture read from the spatial-upsampler checkpoint metadata.""" + + in_channels: int + mid_channels: int + num_blocks_per_stage: int + dims: int + spatial_upsample: bool + temporal_upsample: bool + + @classmethod + def from_metadata(cls, metadata: dict[str, Any]) -> "LTX25SpatialUpsamplerConfig": + config = metadata.get("config") + if not isinstance(config, dict): + raise ValueError("LTX-2.5 spatial upsampler is missing object config metadata") + required = ( + "in_channels", + "mid_channels", + "num_blocks_per_stage", + "dims", + "spatial_upsample", + "temporal_upsample", + ) + missing = [field for field in required if field not in config] + if missing: + raise ValueError(f"LTX-2.5 spatial upsampler config is missing {missing}") + return cls(**{field: config[field] for field in required}) + + +class LTX25SpatialUpsampler(torch.nn.Module): + """Faithful LTX-2.5 learned spatial latent upsampler.""" + + def __init__(self, config: LTX25SpatialUpsamplerConfig) -> None: + super().__init__() + if not config.spatial_upsample or config.temporal_upsample or config.dims not in (2, 3): + raise ValueError("LTX-2.5 distilled requires a spatial-only 2D or 3D latent upsampler") + conv = torch.nn.Conv2d if config.dims == 2 else torch.nn.Conv3d + self.config = config + self.initial_conv = conv(config.in_channels, config.mid_channels, kernel_size=3, padding=1) + self.initial_norm = torch.nn.GroupNorm(32, config.mid_channels) + self.initial_activation = torch.nn.SiLU() + self.res_blocks = torch.nn.ModuleList( + [LTX25UpsamplerResBlock(config.mid_channels, config.dims) for _ in range(config.num_blocks_per_stage)] + ) + self.upsampler = torch.nn.Sequential( + torch.nn.Conv2d(config.mid_channels, 4 * config.mid_channels, kernel_size=3, padding=1), + LTX25PixelShuffleND(2), + ) + self.post_upsample_res_blocks = torch.nn.ModuleList( + [LTX25UpsamplerResBlock(config.mid_channels, config.dims) for _ in range(config.num_blocks_per_stage)] + ) + self.final_conv = conv(config.mid_channels, config.in_channels, kernel_size=3, padding=1) + + def forward(self, latent: torch.Tensor) -> torch.Tensor: + if latent.ndim != 5: + raise ValueError(f"LTX-2.5 latent upsampler expects [B, C, F, H, W], got {tuple(latent.shape)}") + batch, _, frames, _, _ = latent.shape + if self.config.dims == 2: + value = latent.permute(0, 2, 1, 3, 4).reshape( + batch * frames, latent.shape[1], latent.shape[3], latent.shape[4] + ) + value = self.initial_activation(self.initial_norm(self.initial_conv(value))) + for block in self.res_blocks: + value = block(value) + value = self.upsampler(value) + for block in self.post_upsample_res_blocks: + value = block(value) + value = self.final_conv(value) + return value.reshape(batch, frames, value.shape[1], value.shape[2], value.shape[3]).permute(0, 2, 1, 3, 4) + + value = self.initial_activation(self.initial_norm(self.initial_conv(latent))) + for block in self.res_blocks: + value = block(value) + value = value.permute(0, 2, 1, 3, 4).reshape(batch * frames, value.shape[1], value.shape[3], value.shape[4]) + value = self.upsampler(value) + value = value.reshape(batch, frames, value.shape[1], value.shape[2], value.shape[3]).permute(0, 2, 1, 3, 4) + for block in self.post_upsample_res_blocks: + value = block(value) + return self.final_conv(value) + + @classmethod + def from_checkpoint( + cls, + checkpoint_path: str | Path, + *, + device: torch.device | str = "cpu", + torch_dtype: torch.dtype = torch.bfloat16, + ) -> "LTX25SpatialUpsampler": + """Load an isolated LTX-2.5 upsampler with exact checkpoint-key coverage.""" + metadata = inspect_checkpoint(checkpoint_path).metadata + model = cls(LTX25SpatialUpsamplerConfig.from_metadata(metadata)) + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as checkpoint: + state_dict = {key: checkpoint.get_tensor(key) for key in checkpoint.keys()} + missing, unexpected = model.load_state_dict(state_dict, strict=True) + if missing or unexpected: + raise ValueError(f"LTX-2.5 upsampler load mismatch: missing={missing}, unexpected={unexpected}") + return model.to(device=device, dtype=torch_dtype).eval() + + +def load_video_latent_statistics(video_vae_path: str | Path) -> LTX25PerChannelStatistics: + """Load only the VAE normalization buffers needed by the spatial-upsample bridge.""" + source = Path(video_vae_path).expanduser().resolve() + with safe_open(str(source), framework="pt", device="cpu") as checkpoint: + state_dict = { + key.removeprefix("per_channel_statistics."): checkpoint.get_tensor(key) + for key in checkpoint.keys() + if key.startswith("per_channel_statistics.") + } + statistics = LTX25PerChannelStatistics(len(state_dict["std-of-means"])) + missing, unexpected = statistics.load_state_dict(state_dict, strict=True) + if missing or unexpected: + raise ValueError(f"LTX-2.5 video-statistics load mismatch: missing={missing}, unexpected={unexpected}") + return statistics.eval() + + +def upsample_video_latent( + latent: torch.Tensor, + upsampler: LTX25SpatialUpsampler, + statistics: LTX25PerChannelStatistics, +) -> torch.Tensor: + """Unnormalize, spatially upsample, then renormalize an LTX-2.5 video latent.""" + return statistics.normalize(upsampler(statistics.un_normalize(latent))) diff --git a/telefuser/models/ltx25/transformer.py b/telefuser/models/ltx25/transformer.py new file mode 100644 index 0000000..1a215ab --- /dev/null +++ b/telefuser/models/ltx25/transformer.py @@ -0,0 +1,2468 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import torch +from loguru import logger +from safetensors import safe_open +from torch.distributed.device_mesh import DeviceMesh + +from telefuser.core.base_model import BaseModel +from telefuser.core.config import AttentionConfig, AttnImplType, OffloadConfig +from telefuser.distributed import ( + get_attention_strategy, + get_ulysses_group, + get_ulysses_world_size, + ulysses_gather_heads, + ulysses_scatter_heads, +) +from telefuser.distributed.parallel_shard import sequence_parallel_shard, sequence_parallel_unshard +from telefuser.ops.attention import attention as attn_func +from telefuser.ops.normalization import LayerNorm, RMSNorm + + +def rms_norm(x: torch.Tensor, weight: torch.Tensor | None = None, eps: float = 1e-6) -> torch.Tensor: + """Root-mean-square (RMS) normalize `x` over its last dimension. + Thin wrapper around `torch.nn.functional.rms_norm` that infers the normalized + shape and forwards `weight` and `eps`. + """ + return torch.nn.functional.rms_norm(x, (x.shape[-1],), weight=weight, eps=eps) + + +def check_config_value(config: dict, key: str, expected: Any) -> None: # noqa: ANN401 + actual = config.get(key) + if actual != expected: + raise ValueError(f"Config value {key} is {actual}, expected {expected}") + + +def to_velocity( + sample: torch.Tensor, + sigma: float | torch.Tensor, + denoised_sample: torch.Tensor, + calc_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """ + Convert the sample and its denoised version to velocity. + Returns: + Velocity + """ + if isinstance(sigma, torch.Tensor): + sigma = sigma.to(calc_dtype).item() + if sigma == 0: + raise ValueError("Sigma can't be 0.0") + return ((sample.to(calc_dtype) - denoised_sample.to(calc_dtype)) / sigma).to(sample.dtype) + + +def to_denoised( + sample: torch.Tensor, + velocity: torch.Tensor, + sigma: float | torch.Tensor, + calc_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """ + Convert the sample and its denoising velocity to denoised sample. + Returns: + Denoised sample + """ + if isinstance(sigma, torch.Tensor): + sigma = sigma.to(calc_dtype) + return (sample.to(calc_dtype) - velocity.to(calc_dtype) * sigma).to(sample.dtype) + + +def find_matching_file(root_path: str, pattern: str) -> Path: + """ + Recursively search for files matching a glob pattern and return the first match. + """ + matches = list(Path(root_path).rglob(pattern)) + if not matches: + raise FileNotFoundError(f"No files matching pattern '{pattern}' found under {root_path}") + return matches[0] + + +from dataclasses import dataclass +from enum import Enum + +import torch +from torch._prims_common import DeviceLikeType + + +class PerturbationType(Enum): + """Types of attention perturbations for STG (Spatio-Temporal Guidance).""" + + SKIP_A2V_CROSS_ATTN = "skip_a2v_cross_attn" + SKIP_V2A_CROSS_ATTN = "skip_v2a_cross_attn" + SKIP_VIDEO_SELF_ATTN = "skip_video_self_attn" + SKIP_AUDIO_SELF_ATTN = "skip_audio_self_attn" + + +@dataclass(frozen=True) +class Perturbation: + """A single perturbation specifying which attention type to skip and in which blocks.""" + + type: PerturbationType + blocks: list[int] | None # None means all blocks + + def is_perturbed(self, perturbation_type: PerturbationType, block: int) -> bool: + if self.type != perturbation_type: + return False + + if self.blocks is None: + return True + + return block in self.blocks + + +@dataclass(frozen=True) +class PerturbationConfig: + """Configuration holding a list of perturbations for a single sample.""" + + perturbations: list[Perturbation] | None + + def is_perturbed(self, perturbation_type: PerturbationType, block: int) -> bool: + if self.perturbations is None: + return False + + return any(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations) + + @staticmethod + def empty() -> "PerturbationConfig": + return PerturbationConfig([]) + + +@dataclass(frozen=True) +class BatchedPerturbationConfig: + """Perturbation configurations for a batch, with utilities for generating attention masks.""" + + perturbations: list[PerturbationConfig] + + def mask( + self, perturbation_type: PerturbationType, block: int, device: DeviceLikeType, dtype: torch.dtype + ) -> torch.Tensor: + mask = torch.ones((len(self.perturbations),), device=device, dtype=dtype) + for batch_idx, perturbation in enumerate(self.perturbations): + if perturbation.is_perturbed(perturbation_type, block): + mask[batch_idx] = 0 + + return mask + + def mask_like(self, perturbation_type: PerturbationType, block: int, values: torch.Tensor) -> torch.Tensor: + mask = self.mask(perturbation_type, block, values.device, values.dtype) + return mask.view(mask.numel(), *([1] * len(values.shape[1:]))) + + def any_in_batch(self, perturbation_type: PerturbationType, block: int) -> bool: + return any(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations) + + def all_in_batch(self, perturbation_type: PerturbationType, block: int) -> bool: + return all(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations) + + @staticmethod + def empty(batch_size: int) -> "BatchedPerturbationConfig": + return BatchedPerturbationConfig([PerturbationConfig.empty() for _ in range(batch_size)]) + + +import torch + + +class GELUApprox(torch.nn.Module): + def __init__(self, dim_in: int, dim_out: int, bias: bool = True) -> None: + super().__init__() + self.proj = torch.nn.Linear(dim_in, dim_out, bias=bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.gelu(self.proj(x), approximate="tanh") + + +import math + +import torch + + +def get_timestep_embedding( + timesteps: torch.Tensor, + embedding_dim: int, + flip_sin_to_cos: bool = False, + downscale_freq_shift: float = 1, + scale: float = 1, + max_period: int = 10000, +) -> torch.Tensor: + """ + This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. + Args + timesteps (torch.Tensor): + a 1-D Tensor of N indices, one per batch element. These may be fractional. + embedding_dim (int): + the dimension of the output. + flip_sin_to_cos (bool): + Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False) + downscale_freq_shift (float): + Controls the delta between frequencies between dimensions + scale (float): + Scaling factor applied to the embeddings. + max_period (int): + Controls the maximum frequency of the embeddings + Returns + torch.Tensor: an [N x dim] Tensor of positional embeddings. + """ + assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" + + half_dim = embedding_dim // 2 + exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device) + exponent = exponent / (half_dim - downscale_freq_shift) + + emb = torch.exp(exponent) + emb = timesteps[:, None].float() * emb[None, :] + + # scale embeddings + emb = scale * emb + + # concat sine and cosine embeddings + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + + # flip sine and cosine embeddings + if flip_sin_to_cos: + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) + + # zero pad + if embedding_dim % 2 == 1: + emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) + return emb + + +class TimestepEmbedding(torch.nn.Module): + def __init__( + self, + in_channels: int, + time_embed_dim: int, + out_dim: int | None = None, + post_act_fn: str | None = None, + cond_proj_dim: int | None = None, + sample_proj_bias: bool = True, + ): + super().__init__() + + self.linear_1 = torch.nn.Linear(in_channels, time_embed_dim, sample_proj_bias) + + if cond_proj_dim is not None: + self.cond_proj = torch.nn.Linear(cond_proj_dim, in_channels, bias=False) + else: + self.cond_proj = None + + self.act = torch.nn.SiLU() + time_embed_dim_out = out_dim if out_dim is not None else time_embed_dim + + self.linear_2 = torch.nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias) + + if post_act_fn is None: + self.post_act = None + + def forward(self, sample: torch.Tensor, condition: torch.Tensor | None = None) -> torch.Tensor: + if condition is not None: + sample = sample + self.cond_proj(condition) + sample = self.linear_1(sample) + + if self.act is not None: + sample = self.act(sample) + + sample = self.linear_2(sample) + + if self.post_act is not None: + sample = self.post_act(sample) + return sample + + +class Timesteps(torch.nn.Module): + def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1): + super().__init__() + self.num_channels = num_channels + self.flip_sin_to_cos = flip_sin_to_cos + self.downscale_freq_shift = downscale_freq_shift + self.scale = scale + + def forward(self, timesteps: torch.Tensor) -> torch.Tensor: + t_emb = get_timestep_embedding( + timesteps, + self.num_channels, + flip_sin_to_cos=self.flip_sin_to_cos, + downscale_freq_shift=self.downscale_freq_shift, + scale=self.scale, + ) + return t_emb + + +class PixArtAlphaCombinedTimestepSizeEmbeddings(torch.nn.Module): + """ + For PixArt-Alpha. + Reference: + https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L164C9-L168C29 + """ + + def __init__( + self, + embedding_dim: int, + size_emb_dim: int, + ): + super().__init__() + + self.outdim = size_emb_dim + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + def forward( + self, + timestep: torch.Tensor, + hidden_dtype: torch.dtype, + ) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D) + return timesteps_emb + + +from typing import Optional, Tuple + +import torch + +# Number of AdaLN modulation parameters per transformer block. +# Base: 2 params (shift + scale) x 3 norms (self-attn, feed-forward, output). +ADALN_NUM_BASE_PARAMS = 6 +# Cross-attention AdaLN adds 3 more (scale, shift, gate) for the CA norm. +ADALN_NUM_CROSS_ATTN_PARAMS = 3 + + +def adaln_embedding_coefficient(cross_attention_adaln: bool) -> int: + """Total number of AdaLN parameters per block.""" + return ADALN_NUM_BASE_PARAMS + (ADALN_NUM_CROSS_ATTN_PARAMS if cross_attention_adaln else 0) + + +class AdaLayerNormSingle(torch.nn.Module): + r""" + Norm layer adaptive layer norm single (adaLN-single). + As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3). + Parameters: + embedding_dim (`int`): The size of each embedding vector. + use_additional_conditions (`bool`): To use additional conditions for normalization or not. + """ + + def __init__(self, embedding_dim: int, embedding_coefficient: int = 6): + super().__init__() + + self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings( + embedding_dim, + size_emb_dim=embedding_dim // 3, + ) + + self.silu = torch.nn.SiLU() + self.linear = torch.nn.Linear(embedding_dim, embedding_coefficient * embedding_dim, bias=True) + + def forward( + self, + timestep: torch.Tensor, + hidden_dtype: Optional[torch.dtype] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype) + return self.linear(self.silu(embedded_timestep)), embedded_timestep + + +import functools +from enum import Enum +from typing import Callable + +import numpy as np +import torch +from einops import rearrange + + +class LTXRopeType(Enum): + INTERLEAVED = "interleaved" + SPLIT = "split" + + +def apply_rotary_emb( + input_tensor: torch.Tensor, + freqs_cis: Tuple[torch.Tensor, torch.Tensor], + rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, +) -> torch.Tensor: + if rope_type == LTXRopeType.INTERLEAVED: + return apply_interleaved_rotary_emb(input_tensor, *freqs_cis) + elif rope_type == LTXRopeType.SPLIT: + return apply_split_rotary_emb(input_tensor, *freqs_cis) + else: + raise ValueError(f"Invalid rope type: {rope_type}") + + +def apply_interleaved_rotary_emb( + input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor +) -> torch.Tensor: + t_dup = rearrange(input_tensor, "... (d r) -> ... d r", r=2) + t1, t2 = t_dup.unbind(dim=-1) + t_dup = torch.stack((-t2, t1), dim=-1) + input_tensor_rot = rearrange(t_dup, "... d r -> ... (d r)") + + out = input_tensor * cos_freqs + input_tensor_rot * sin_freqs + + return out + + +def apply_split_rotary_emb( + input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor +) -> torch.Tensor: + needs_reshape = False + if input_tensor.ndim != 4 and cos_freqs.ndim == 4: + b, h, t, _ = cos_freqs.shape + input_tensor = input_tensor.reshape(b, t, h, -1).swapaxes(1, 2) + needs_reshape = True + + split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2) + first_half_input = split_input[..., :1, :] + second_half_input = split_input[..., 1:, :] + + output = split_input * cos_freqs.unsqueeze(-2) + first_half_output = output[..., :1, :] + second_half_output = output[..., 1:, :] + + first_half_output.addcmul_(-sin_freqs.unsqueeze(-2), second_half_input) + second_half_output.addcmul_(sin_freqs.unsqueeze(-2), first_half_input) + + output = rearrange(output, "... d r -> ... (d r)") + if needs_reshape: + output = output.swapaxes(1, 2).reshape(b, t, -1) + + return output + + +@functools.lru_cache(maxsize=5) +def generate_freq_grid_np( + positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int +) -> torch.Tensor: + theta = positional_embedding_theta + start = 1 + end = theta + + n_elem = 2 * positional_embedding_max_pos_count + pow_indices = np.power( + theta, + np.linspace( + np.log(start) / np.log(theta), + np.log(end) / np.log(theta), + inner_dim // n_elem, + dtype=np.float64, + ), + ) + return torch.tensor(pow_indices * math.pi / 2, dtype=torch.float32) + + +@functools.lru_cache(maxsize=5) +def generate_freq_grid_pytorch( + positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int +) -> torch.Tensor: + theta = positional_embedding_theta + start = 1 + end = theta + n_elem = 2 * positional_embedding_max_pos_count + + indices = theta ** ( + torch.linspace( + math.log(start, theta), + math.log(end, theta), + inner_dim // n_elem, + dtype=torch.float32, + ) + ) + indices = indices.to(dtype=torch.float32) + + indices = indices * math.pi / 2 + + return indices + + +def get_fractional_positions(indices_grid: torch.Tensor, max_pos: list[int]) -> torch.Tensor: + n_pos_dims = indices_grid.shape[1] + assert n_pos_dims == len(max_pos), ( + f"Number of position dimensions ({n_pos_dims}) must match max_pos length ({len(max_pos)})" + ) + fractional_positions = torch.stack( + [indices_grid[:, i] / max_pos[i] for i in range(n_pos_dims)], + dim=-1, + ) + return fractional_positions + + +def generate_freqs( + indices: torch.Tensor, indices_grid: torch.Tensor, max_pos: list[int], use_middle_indices_grid: bool +) -> torch.Tensor: + if use_middle_indices_grid: + if len(indices_grid.shape) != 4: + raise ValueError( + "use_middle_indices_grid expects indices_grid with shape (B, D, T, bounds), got " + f"{tuple(indices_grid.shape)}." + ) + # Video positions include [start, end] bounds for each token; audio positions may only provide a single + # coordinate (equivalent to start) which is sufficient when patch size is 1. + if indices_grid.shape[-1] == 2: + indices_grid_start, indices_grid_end = indices_grid[..., 0], indices_grid[..., 1] + indices_grid = (indices_grid_start + indices_grid_end) / 2.0 + elif indices_grid.shape[-1] == 1: + indices_grid = indices_grid[..., 0] + else: + raise ValueError( + "use_middle_indices_grid expects indices_grid bounds dimension to be 1 or 2, got " + f"{indices_grid.shape[-1]} for shape {tuple(indices_grid.shape)}." + ) + elif len(indices_grid.shape) == 4: + indices_grid = indices_grid[..., 0] + + fractional_positions = get_fractional_positions(indices_grid, max_pos) + indices = indices.to(device=fractional_positions.device) + + freqs = (indices * (fractional_positions.unsqueeze(-1) * 2 - 1)).transpose(-1, -2).flatten(2) + return freqs + + +def split_freqs_cis(freqs: torch.Tensor, pad_size: int, num_attention_heads: int) -> tuple[torch.Tensor, torch.Tensor]: + cos_freq = freqs.cos() + sin_freq = freqs.sin() + + if pad_size != 0: + cos_padding = torch.ones_like(cos_freq[:, :, :pad_size]) + sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size]) + + cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1) + sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1) + + # Reshape freqs to be compatible with multi-head attention + b = cos_freq.shape[0] + t = cos_freq.shape[1] + + cos_freq = cos_freq.reshape(b, t, num_attention_heads, -1) + sin_freq = sin_freq.reshape(b, t, num_attention_heads, -1) + + cos_freq = torch.swapaxes(cos_freq, 1, 2) # (B,H,T,D//2) + sin_freq = torch.swapaxes(sin_freq, 1, 2) # (B,H,T,D//2) + return cos_freq, sin_freq + + +def interleaved_freqs_cis(freqs: torch.Tensor, pad_size: int) -> tuple[torch.Tensor, torch.Tensor]: + cos_freq = freqs.cos().repeat_interleave(2, dim=-1) + sin_freq = freqs.sin().repeat_interleave(2, dim=-1) + if pad_size != 0: + cos_padding = torch.ones_like(cos_freq[:, :, :pad_size]) + sin_padding = torch.zeros_like(cos_freq[:, :, :pad_size]) + cos_freq = torch.cat([cos_padding, cos_freq], dim=-1) + sin_freq = torch.cat([sin_padding, sin_freq], dim=-1) + return cos_freq, sin_freq + + +def precompute_freqs_cis( + indices_grid: torch.Tensor, + dim: int, + out_dtype: torch.dtype, + theta: float = 10000.0, + max_pos: list[int] | None = None, + use_middle_indices_grid: bool = False, + num_attention_heads: int = 32, + rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, + freq_grid_generator: Callable[[float, int, int, torch.device], torch.Tensor] = generate_freq_grid_pytorch, +) -> tuple[torch.Tensor, torch.Tensor]: + if max_pos is None: + max_pos = [20, 2048, 2048] + + indices = freq_grid_generator(theta, indices_grid.shape[1], dim) + freqs = generate_freqs(indices, indices_grid, max_pos, use_middle_indices_grid) + + if rope_type == LTXRopeType.SPLIT: + expected_freqs = dim // 2 + current_freqs = freqs.shape[-1] + pad_size = expected_freqs - current_freqs + cos_freq, sin_freq = split_freqs_cis(freqs, pad_size, num_attention_heads) + else: + # 2 because of cos and sin by 3 for (t, x, y), 1 for temporal only + n_elem = 2 * indices_grid.shape[1] + cos_freq, sin_freq = interleaved_freqs_cis(freqs, dim % n_elem) + return cos_freq.to(out_dtype), sin_freq.to(out_dtype) + + +from enum import Enum + + +class Attention(torch.nn.Module): + """Multi-head attention that delegates computation to `telefuser.ops.attention`. + + The attention backend is configured globally via `Attention.attention_config`, which is intended to be set by + the runtime (e.g. `ModelRuntimeConfig.attention_config`) like other TeleFuser models. + """ + + attention_config: AttentionConfig = AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA) + + def __init__( + self, + query_dim: int, + context_dim: int | None = None, + heads: int = 8, + dim_head: int = 64, + norm_eps: float = 1e-6, + rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, + apply_gated_attention: bool = False, + ) -> None: + super().__init__() + self.rope_type = rope_type + + inner_dim = dim_head * heads + context_dim = query_dim if context_dim is None else context_dim + + self.heads = heads + self.dim_head = dim_head + self.ulysses_group: torch.distributed.ProcessGroup | None = None + + # The LTX-2.5 reference uses PyTorch RMSNorm. Keep that exact numerical + # path for the fidelity baseline; attention itself still goes through the + # public TeleFuser dispatch layer below. + self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) + self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) + + self.to_q = torch.nn.Linear(query_dim, inner_dim, bias=True) + self.to_k = torch.nn.Linear(context_dim, inner_dim, bias=True) + self.to_v = torch.nn.Linear(context_dim, inner_dim, bias=True) + + # Optional per-head gating + if apply_gated_attention: + self.to_gate_logits = torch.nn.Linear(query_dim, heads, bias=True) + else: + self.to_gate_logits = None + + self.to_out = torch.nn.Sequential(torch.nn.Linear(inner_dim, query_dim, bias=True), torch.nn.Identity()) + + def set_ulysses_group(self, process_group: torch.distributed.ProcessGroup | None) -> None: + """Configure the process group used for Ulysses head/sequence exchange.""" + self.ulysses_group = process_group + + def forward( + self, + x: torch.Tensor, + context: torch.Tensor | None = None, + mask: torch.Tensor | None = None, + pe: torch.Tensor | None = None, + k_pe: torch.Tensor | None = None, + perturbation_mask: torch.Tensor | None = None, + all_perturbed: bool = False, + enforce_mask: bool = False, + ) -> torch.Tensor: + """Multi-head attention with optional RoPE, perturbation masking, and per-head gating. + When ``perturbation_mask`` is all zeros, the expensive query/key path + (linear projections, RMSNorm, RoPE) is skipped entirely and only the + value projection is used as a pass-through. + Args: + x: Query input tensor of shape ``(B, T, query_dim)``. + context: Key/value context tensor of shape ``(B, S, context_dim)``. + Falls back to ``x`` (self-attention) when *None*. + mask: Optional attention mask. Interpretation depends on the attention + backend (additive bias for xformers/PyTorch SDPA). + pe: Rotary positional embeddings applied to both ``q`` and ``k``. + k_pe: Separate rotary positional embeddings for ``k`` only. When + *None*, ``pe`` is reused for keys. + perturbation_mask: Optional mask in ``[0, 1]`` that + blends the attention output with the raw value projection: + ``out = attn_out * mask + v * (1 - mask)``. + **1** keeps the full attention output, **0** bypasses attention + and passes the value projection through unchanged. + *None* or all-ones means standard attention; all-zeros skips + the query/key path entirely for efficiency. + all_perturbed: Whether all perturbations are active for this block. + enforce_mask: Route this call through SDPA so padding masks are honored. + Returns: + Output tensor of shape ``(B, T, query_dim)``. + """ + context = x if context is None else context + use_attention = not all_perturbed + + v = self.to_v(context) + + if not use_attention: + out = v + else: + q = self.to_q(x) + k = self.to_k(context) + + q = self.q_norm(q) + k = self.k_norm(k) + + if pe is not None: + q = apply_rotary_emb(q, pe, self.rope_type) + k = apply_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type) + + q = rearrange(q, "b t (h d) -> b t h d", h=self.heads) + k = rearrange(k, "b s (h d) -> b s h d", h=self.heads) + v = rearrange(v, "b s (h d) -> b s h d", h=self.heads) + local_v = v + if self.ulysses_group is not None: + q_wait = ulysses_scatter_heads(q, self.ulysses_group) + k_wait = ulysses_scatter_heads(k, self.ulysses_group) + v_wait = ulysses_scatter_heads(v, self.ulysses_group) + q, k, v = q_wait(), k_wait(), v_wait() + attention_config = self.attention_config + if enforce_mask: + attention_config = AttentionConfig( + attn_impl=AttnImplType.TORCH_SDPA, + scale=attention_config.scale, + dropout=attention_config.dropout, + is_causal=attention_config.is_causal, + ) + out = attn_func(q, k, v, attention_config=attention_config, attn_mask=mask, input_layout="BSND") + if self.ulysses_group is not None: + out = ulysses_gather_heads(out, self.ulysses_group, num_heads=self.heads)() + out = rearrange(out, "b t h d -> b t (h d)") + + if perturbation_mask is not None: + out = out * perturbation_mask + rearrange(local_v, "b s h d -> b s (h d)") * (1 - perturbation_mask) + + # Apply per-head gating if enabled + if self.to_gate_logits is not None: + gate_logits = self.to_gate_logits(x) # (B, T, H) + b, t, _ = out.shape + # Reshape to (B, T, H, D) for per-head gating + out = out.view(b, t, self.heads, self.dim_head) + # Apply gating: 2 * sigmoid(x) so that zero-init gives identity (2 * 0.5 = 1.0) + gates = 2.0 * torch.sigmoid(gate_logits) # (B, T, H) + out = out * gates.unsqueeze(-1) # (B, T, H, D) * (B, T, H, 1) + # Reshape back to (B, T, H*D) + out = out.view(b, t, self.heads * self.dim_head) + + return self.to_out(out) + + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class Modality: + """ + Input data for a single modality (video or audio) in the transformer. + Bundles the latent tokens, timestep embeddings, positional information, + and text conditioning context for processing by the diffusion transformer. + Attributes: + latent: Patchified latent tokens, shape ``(B, T, D)`` where *B* is + the batch size, *T* is the total number of tokens (noisy + + conditioning), and *D* is the input dimension. + timesteps: Per-token timestep embeddings, shape ``(B, T)``. + positions: Positional coordinates, shape ``(B, 3, T)`` for video + (time, height, width) or ``(B, 1, T)`` for audio. + context: Text conditioning embeddings from the prompt encoder. + enabled: Whether this modality is active in the current forward pass. + context_mask: Optional mask for the text context tokens. + attention_mask: Optional 2-D self-attention mask, shape ``(B, T, T)``. + Values in ``[0, 1]`` where ``1`` = full attention and ``0`` = no + attention. ``None`` means unrestricted (full) attention between + all tokens. Built incrementally by conditioning items; see + conditioning strength wrappers. + """ + + latent: ( + torch.Tensor + ) # Shape: (B, T, D) where B is the batch size, T is the number of tokens, and D is input dimension + sigma: torch.Tensor # Shape: (B,). Current sigma value, used for cross-attention timestep calculation. + timesteps: torch.Tensor # Shape: (B, T) where T is the number of timesteps + positions: ( + torch.Tensor + ) # Shape: (B, 3, T) for video, where 3 is the number of dimensions and T is the number of tokens + context: torch.Tensor + enabled: bool = True + context_mask: torch.Tensor | None = None + attention_mask: torch.Tensor | None = None + keyframes_mask: torch.Tensor | None = None + + +from dataclasses import dataclass, replace + +import torch + + +@dataclass(frozen=True) +class TransformerArgs: + x: torch.Tensor + context: torch.Tensor + context_mask: torch.Tensor + timesteps: torch.Tensor + embedded_timestep: torch.Tensor + positional_embeddings: torch.Tensor + cross_positional_embeddings: torch.Tensor | None + cross_scale_shift_timestep: torch.Tensor | None + cross_gate_timestep: torch.Tensor | None + enabled: bool + prompt_timestep: torch.Tensor | None = None + self_attention_mask: torch.Tensor | None = ( + None # Additive log-space self-attention bias (B, 1, T, T), None = full attention + ) + key_padding_mask: torch.Tensor | None = None + + +class TransformerArgsPreprocessor: + def __init__( # noqa: PLR0913 + self, + patchify_proj: torch.nn.Linear, + adaln: AdaLayerNormSingle, + inner_dim: int, + max_pos: list[int], + num_attention_heads: int, + use_middle_indices_grid: bool, + timestep_scale_multiplier: int, + double_precision_rope: bool, + positional_embedding_theta: float, + rope_type: LTXRopeType, + caption_projection: torch.nn.Module | None = None, + prompt_adaln: AdaLayerNormSingle | None = None, + keyframes_embedding_provider: Callable[[], torch.Tensor | None] | None = None, + ) -> None: + self.patchify_proj = patchify_proj + self.adaln = adaln + self.inner_dim = inner_dim + self.max_pos = max_pos + self.num_attention_heads = num_attention_heads + self.use_middle_indices_grid = use_middle_indices_grid + self.timestep_scale_multiplier = timestep_scale_multiplier + self.double_precision_rope = double_precision_rope + self.positional_embedding_theta = positional_embedding_theta + self.rope_type = rope_type + self.caption_projection = caption_projection + self.prompt_adaln = prompt_adaln + self.keyframes_embedding_provider = keyframes_embedding_provider + + def _prepare_timestep( + self, timestep: torch.Tensor, adaln: AdaLayerNormSingle, batch_size: int, hidden_dtype: torch.dtype + ) -> tuple[torch.Tensor, torch.Tensor]: + """Prepare timestep embeddings.""" + timestep_scaled = timestep * self.timestep_scale_multiplier + timestep, embedded_timestep = adaln( + timestep_scaled.flatten(), + hidden_dtype=hidden_dtype, + ) + # Second dimension is 1 or number of tokens (if timestep_per_token) + timestep = timestep.view(batch_size, -1, timestep.shape[-1]) + embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.shape[-1]) + + return timestep, embedded_timestep + + def _prepare_context( + self, + context: torch.Tensor, + x: torch.Tensor, + ) -> torch.Tensor: + """Prepare context for transformer blocks.""" + if self.caption_projection is not None: + context = self.caption_projection(context) + batch_size = x.shape[0] + return context.view(batch_size, -1, x.shape[-1]) + + def _prepare_attention_mask(self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype) -> torch.Tensor | None: + """Prepare attention mask.""" + if attention_mask is None or torch.is_floating_point(attention_mask): + return attention_mask + + return (attention_mask - 1).to(x_dtype).reshape( + (attention_mask.shape[0], 1, -1, attention_mask.shape[-1]) + ) * torch.finfo(x_dtype).max + + def _prepare_self_attention_mask( + self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype + ) -> torch.Tensor | None: + """Prepare self-attention mask by converting [0,1] values to additive log-space bias. + Input shape: (B, T, T) with values in [0, 1]. + Output shape: (B, 1, T, T) with 0.0 for full attention and a large negative value + for masked positions. + Positions with attention_mask <= 0 are fully masked (mapped to the dtype's minimum + representable value). Strictly positive entries are converted via log-space for + smooth attenuation, with small values clamped for numerical stability. + Returns None if input is None (no masking). + """ + if attention_mask is None: + return None + + # Convert [0, 1] attention mask to additive log-space bias: + # 1.0 -> log(1.0) = 0.0 (no bias, full attention) + # 0.0 -> finfo.min (fully masked) + finfo = torch.finfo(x_dtype) + eps = finfo.tiny + + bias = torch.full_like(attention_mask, finfo.min, dtype=x_dtype) + positive = attention_mask > 0 + if positive.any(): + bias[positive] = torch.log(attention_mask[positive].clamp(min=eps)).to(x_dtype) + + return bias.unsqueeze(1) # (B, 1, T, T) for head broadcast + + def _prepare_positional_embeddings( + self, + positions: torch.Tensor, + inner_dim: int, + max_pos: list[int], + use_middle_indices_grid: bool, + num_attention_heads: int, + x_dtype: torch.dtype, + ) -> torch.Tensor: + """Prepare positional embeddings.""" + freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch + pe = precompute_freqs_cis( + positions, + dim=inner_dim, + out_dtype=x_dtype, + theta=self.positional_embedding_theta, + max_pos=max_pos, + use_middle_indices_grid=use_middle_indices_grid, + num_attention_heads=num_attention_heads, + rope_type=self.rope_type, + freq_grid_generator=freq_grid_generator, + ) + return pe + + def prepare( + self, + modality: Modality, + cross_modality: Modality | None = None, # noqa: ARG002 + ) -> TransformerArgs: + x = self.patchify_proj(modality.latent) + if self.keyframes_embedding_provider is not None and modality.keyframes_mask is not None: + keyframes_embedding = self.keyframes_embedding_provider() + if keyframes_embedding is not None: + x = x + (modality.keyframes_mask > 0).to(dtype=x.dtype) * keyframes_embedding.to(dtype=x.dtype) + batch_size = x.shape[0] + timestep, embedded_timestep = self._prepare_timestep( + modality.timesteps, self.adaln, batch_size, modality.latent.dtype + ) + prompt_timestep = None + if self.prompt_adaln is not None: + prompt_timestep, _ = self._prepare_timestep( + modality.sigma, self.prompt_adaln, batch_size, modality.latent.dtype + ) + context = self._prepare_context(modality.context, x) + attention_mask = self._prepare_attention_mask(modality.context_mask, modality.latent.dtype) + pe = self._prepare_positional_embeddings( + positions=modality.positions, + inner_dim=self.inner_dim, + max_pos=self.max_pos, + use_middle_indices_grid=self.use_middle_indices_grid, + num_attention_heads=self.num_attention_heads, + x_dtype=modality.latent.dtype, + ) + self_attention_mask = self._prepare_self_attention_mask(modality.attention_mask, modality.latent.dtype) + return TransformerArgs( + x=x, + context=context, + context_mask=attention_mask, + timesteps=timestep, + embedded_timestep=embedded_timestep, + positional_embeddings=pe, + cross_positional_embeddings=None, + cross_scale_shift_timestep=None, + cross_gate_timestep=None, + enabled=modality.enabled, + prompt_timestep=prompt_timestep, + self_attention_mask=self_attention_mask, + ) + + +class MultiModalTransformerArgsPreprocessor: + def __init__( # noqa: PLR0913 + self, + patchify_proj: torch.nn.Linear, + adaln: AdaLayerNormSingle, + cross_scale_shift_adaln: AdaLayerNormSingle, + cross_gate_adaln: AdaLayerNormSingle, + inner_dim: int, + max_pos: list[int], + num_attention_heads: int, + cross_pe_max_pos: int, + use_middle_indices_grid: bool, + audio_cross_attention_dim: int, + timestep_scale_multiplier: int, + double_precision_rope: bool, + positional_embedding_theta: float, + rope_type: LTXRopeType, + av_ca_timestep_scale_multiplier: int, + caption_projection: torch.nn.Module | None = None, + prompt_adaln: AdaLayerNormSingle | None = None, + keyframes_embedding_provider: Callable[[], torch.Tensor | None] | None = None, + ) -> None: + self.simple_preprocessor = TransformerArgsPreprocessor( + patchify_proj=patchify_proj, + adaln=adaln, + inner_dim=inner_dim, + max_pos=max_pos, + num_attention_heads=num_attention_heads, + use_middle_indices_grid=use_middle_indices_grid, + timestep_scale_multiplier=timestep_scale_multiplier, + double_precision_rope=double_precision_rope, + positional_embedding_theta=positional_embedding_theta, + rope_type=rope_type, + caption_projection=caption_projection, + prompt_adaln=prompt_adaln, + keyframes_embedding_provider=keyframes_embedding_provider, + ) + self.cross_scale_shift_adaln = cross_scale_shift_adaln + self.cross_gate_adaln = cross_gate_adaln + self.cross_pe_max_pos = cross_pe_max_pos + self.audio_cross_attention_dim = audio_cross_attention_dim + self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier + + def prepare( + self, + modality: Modality, + cross_modality: Modality | None = None, + ) -> TransformerArgs: + transformer_args = self.simple_preprocessor.prepare(modality) + if cross_modality is None: + return transformer_args + + if cross_modality.sigma.numel() > 1: + if cross_modality.sigma.shape[0] != modality.timesteps.shape[0]: + raise ValueError("Cross modality sigma must have the same batch size as the modality") + if cross_modality.sigma.ndim != 1: + raise ValueError("Cross modality sigma must be a 1D tensor") + + cross_pe = self.simple_preprocessor._prepare_positional_embeddings( + positions=modality.positions[:, 0:1, :], + inner_dim=self.audio_cross_attention_dim, + max_pos=[self.cross_pe_max_pos], + use_middle_indices_grid=True, + num_attention_heads=self.simple_preprocessor.num_attention_heads, + x_dtype=modality.latent.dtype, + ) + + cross_scale_shift_timestep, cross_gate_timestep = self._prepare_cross_attention_timestep( + modality_timesteps=modality.timesteps, + cross_modality_sigma=cross_modality.sigma, + timestep_scale_multiplier=self.simple_preprocessor.timestep_scale_multiplier, + batch_size=transformer_args.x.shape[0], + hidden_dtype=modality.latent.dtype, + ) + + return replace( + transformer_args, + cross_positional_embeddings=cross_pe, + cross_scale_shift_timestep=cross_scale_shift_timestep, + cross_gate_timestep=cross_gate_timestep, + ) + + def _prepare_cross_attention_timestep( + self, + modality_timesteps: torch.Tensor, + cross_modality_sigma: torch.Tensor, + timestep_scale_multiplier: int, + batch_size: int, + hidden_dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Prepare cross-attention AdaLN inputs with upstream token granularity.""" + av_ca_factor = self.av_ca_timestep_scale_multiplier / timestep_scale_multiplier + scale_shift_timestep, _ = self.cross_scale_shift_adaln( + (modality_timesteps * timestep_scale_multiplier).flatten(), + hidden_dtype=hidden_dtype, + ) + scale_shift_timestep = scale_shift_timestep.view(batch_size, -1, scale_shift_timestep.shape[-1]) + gate_noise_timestep, _ = self.cross_gate_adaln( + (cross_modality_sigma * timestep_scale_multiplier * av_ca_factor).flatten(), + hidden_dtype=hidden_dtype, + ) + gate_noise_timestep = gate_noise_timestep.view(batch_size, -1, gate_noise_timestep.shape[-1]) + + return scale_shift_timestep, gate_noise_timestep + + +import torch + + +class FeedForward(torch.nn.Module): + def __init__(self, dim: int, dim_out: int, mult: int = 4, bias: bool = True) -> None: + super().__init__() + inner_dim = int(dim * mult) + project_in = GELUApprox(dim, inner_dim, bias=bias) + + self.net = torch.nn.Sequential(project_in, torch.nn.Identity(), torch.nn.Linear(inner_dim, dim_out, bias=bias)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + +from dataclasses import dataclass + +import torch + + +@dataclass +class TransformerConfig: + dim: int + heads: int + d_head: int + context_dim: int + apply_gated_attention: bool = False + cross_attention_adaln: bool = False + ff_bias: bool = True + + +class BasicAVTransformerBlock(torch.nn.Module): + def __init__( + self, + idx: int, + video: TransformerConfig | None = None, + audio: TransformerConfig | None = None, + rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, + norm_eps: float = 1e-6, + ): + super().__init__() + + self.idx = idx + if video is not None: + self.attn1 = Attention( + query_dim=video.dim, + heads=video.heads, + dim_head=video.d_head, + context_dim=None, + rope_type=rope_type, + norm_eps=norm_eps, + apply_gated_attention=video.apply_gated_attention, + ) + self.attn2 = Attention( + query_dim=video.dim, + context_dim=video.context_dim, + heads=video.heads, + dim_head=video.d_head, + rope_type=rope_type, + norm_eps=norm_eps, + apply_gated_attention=video.apply_gated_attention, + ) + self.ff = FeedForward(video.dim, dim_out=video.dim, bias=video.ff_bias) + video_sst_size = adaln_embedding_coefficient(video.cross_attention_adaln) + self.scale_shift_table = torch.nn.Parameter(torch.empty(video_sst_size, video.dim)) + + if audio is not None: + self.audio_attn1 = Attention( + query_dim=audio.dim, + heads=audio.heads, + dim_head=audio.d_head, + context_dim=None, + rope_type=rope_type, + norm_eps=norm_eps, + apply_gated_attention=audio.apply_gated_attention, + ) + self.audio_attn2 = Attention( + query_dim=audio.dim, + context_dim=audio.context_dim, + heads=audio.heads, + dim_head=audio.d_head, + rope_type=rope_type, + norm_eps=norm_eps, + apply_gated_attention=audio.apply_gated_attention, + ) + self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim, bias=audio.ff_bias) + audio_sst_size = adaln_embedding_coefficient(audio.cross_attention_adaln) + self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(audio_sst_size, audio.dim)) + + if audio is not None and video is not None: + # Q: Video, K,V: Audio + self.audio_to_video_attn = Attention( + query_dim=video.dim, + context_dim=audio.dim, + heads=audio.heads, + dim_head=audio.d_head, + rope_type=rope_type, + norm_eps=norm_eps, + apply_gated_attention=video.apply_gated_attention, + ) + + # Q: Audio, K,V: Video + self.video_to_audio_attn = Attention( + query_dim=audio.dim, + context_dim=video.dim, + heads=audio.heads, + dim_head=audio.d_head, + rope_type=rope_type, + norm_eps=norm_eps, + apply_gated_attention=audio.apply_gated_attention, + ) + + self.scale_shift_table_a2v_ca_audio = torch.nn.Parameter(torch.empty(5, audio.dim)) + self.scale_shift_table_a2v_ca_video = torch.nn.Parameter(torch.empty(5, video.dim)) + + self.cross_attention_adaln = (video is not None and video.cross_attention_adaln) or ( + audio is not None and audio.cross_attention_adaln + ) + + if self.cross_attention_adaln and video is not None: + self.prompt_scale_shift_table = torch.nn.Parameter(torch.empty(2, video.dim)) + if self.cross_attention_adaln and audio is not None: + self.audio_prompt_scale_shift_table = torch.nn.Parameter(torch.empty(2, audio.dim)) + + self.norm_eps = norm_eps + + def get_ada_values( + self, scale_shift_table: torch.Tensor, batch_size: int, timestep: torch.Tensor, indices: slice + ) -> tuple[torch.Tensor, ...]: + num_ada_params = scale_shift_table.shape[0] + + ada_values = ( + scale_shift_table[indices].unsqueeze(0).unsqueeze(0).to(device=timestep.device, dtype=timestep.dtype) + + timestep.reshape(batch_size, timestep.shape[1], num_ada_params, -1)[:, :, indices, :] + ).unbind(dim=2) + return ada_values + + def get_av_ca_ada_values( + self, + scale_shift_table: torch.Tensor, + batch_size: int, + scale_shift_timestep: torch.Tensor, + gate_timestep: torch.Tensor, + scale_shift_indices: slice, + num_scale_shift_values: int = 4, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + scale_shift_ada_values = self.get_ada_values( + scale_shift_table[:num_scale_shift_values, :], batch_size, scale_shift_timestep, scale_shift_indices + ) + gate_ada_values = self.get_ada_values( + scale_shift_table[num_scale_shift_values:, :], batch_size, gate_timestep, slice(None, None) + ) + + scale, shift = (t.squeeze(2) for t in scale_shift_ada_values) + (gate,) = (t.squeeze(2) for t in gate_ada_values) + + return scale, shift, gate + + def _apply_text_cross_attention( + self, + x: torch.Tensor, + context: torch.Tensor, + attn: Attention, + scale_shift_table: torch.Tensor, + prompt_scale_shift_table: torch.Tensor | None, + timestep: torch.Tensor, + prompt_timestep: torch.Tensor | None, + context_mask: torch.Tensor | None, + cross_attention_adaln: bool = False, + ) -> torch.Tensor: + """Apply text cross-attention, with optional AdaLN modulation.""" + if cross_attention_adaln: + shift_q, scale_q, gate = self.get_ada_values(scale_shift_table, x.shape[0], timestep, slice(6, 9)) + return apply_cross_attention_adaln( + x, + context, + attn, + shift_q, + scale_q, + gate, + prompt_scale_shift_table, + prompt_timestep, + context_mask, + self.norm_eps, + ) + return attn(rms_norm(x, eps=self.norm_eps), context=context, mask=context_mask) + + def forward( # noqa: PLR0915 + self, + video: TransformerArgs | None, + audio: TransformerArgs | None, + perturbations: BatchedPerturbationConfig | None = None, + ) -> tuple[TransformerArgs | None, TransformerArgs | None]: + if video is None and audio is None: + raise ValueError("At least one of video or audio must be provided") + + batch_size = (video or audio).x.shape[0] + + if perturbations is None: + perturbations = BatchedPerturbationConfig.empty(batch_size) + + vx = video.x if video is not None else None + ax = audio.x if audio is not None else None + + run_vx = video is not None and video.enabled and vx.numel() > 0 + run_ax = audio is not None and audio.enabled and ax.numel() > 0 + + run_a2v = run_vx and (audio is not None and ax.numel() > 0) + run_v2a = run_ax and (video is not None and vx.numel() > 0) + + if run_vx: + vshift_msa, vscale_msa, vgate_msa = self.get_ada_values( + self.scale_shift_table, vx.shape[0], video.timesteps, slice(0, 3) + ) + norm_vx = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa + del vshift_msa, vscale_msa + + all_perturbed = perturbations.all_in_batch(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx) + none_perturbed = not perturbations.any_in_batch(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx) + v_mask = ( + perturbations.mask_like(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx, vx) + if not all_perturbed and not none_perturbed + else None + ) + vx = ( + vx + + self.attn1( + norm_vx, + pe=video.positional_embeddings, + mask=video.self_attention_mask, + perturbation_mask=v_mask, + all_perturbed=all_perturbed, + enforce_mask=video.key_padding_mask is not None, + ) + * vgate_msa + ) + del vgate_msa, norm_vx, v_mask + vx = vx + self._apply_text_cross_attention( + vx, + video.context, + self.attn2, + self.scale_shift_table, + getattr(self, "prompt_scale_shift_table", None), + video.timesteps, + video.prompt_timestep, + video.context_mask, + cross_attention_adaln=self.cross_attention_adaln, + ) + + if run_ax: + ashift_msa, ascale_msa, agate_msa = self.get_ada_values( + self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(0, 3) + ) + + norm_ax = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa + del ashift_msa, ascale_msa + all_perturbed = perturbations.all_in_batch(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx) + none_perturbed = not perturbations.any_in_batch(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx) + a_mask = ( + perturbations.mask_like(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx, ax) + if not all_perturbed and not none_perturbed + else None + ) + ax = ( + ax + + self.audio_attn1( + norm_ax, + pe=audio.positional_embeddings, + mask=audio.self_attention_mask, + perturbation_mask=a_mask, + all_perturbed=all_perturbed, + enforce_mask=audio.key_padding_mask is not None, + ) + * agate_msa + ) + del agate_msa, norm_ax, a_mask + ax = ax + self._apply_text_cross_attention( + ax, + audio.context, + self.audio_attn2, + self.audio_scale_shift_table, + getattr(self, "audio_prompt_scale_shift_table", None), + audio.timesteps, + audio.prompt_timestep, + audio.context_mask, + cross_attention_adaln=self.cross_attention_adaln, + ) + + # Audio - Video cross attention. + if run_a2v or run_v2a: + vx_norm3 = rms_norm(vx, eps=self.norm_eps) + ax_norm3 = rms_norm(ax, eps=self.norm_eps) + + if run_a2v and not perturbations.all_in_batch(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx): + scale_ca_video_a2v, shift_ca_video_a2v, gate_out_a2v = self.get_av_ca_ada_values( + self.scale_shift_table_a2v_ca_video, + vx.shape[0], + video.cross_scale_shift_timestep, + video.cross_gate_timestep, + slice(0, 2), + ) + vx_scaled = vx_norm3 * (1 + scale_ca_video_a2v) + shift_ca_video_a2v + del scale_ca_video_a2v, shift_ca_video_a2v + + scale_ca_audio_a2v, shift_ca_audio_a2v, _ = self.get_av_ca_ada_values( + self.scale_shift_table_a2v_ca_audio, + ax.shape[0], + audio.cross_scale_shift_timestep, + audio.cross_gate_timestep, + slice(0, 2), + ) + ax_scaled = ax_norm3 * (1 + scale_ca_audio_a2v) + shift_ca_audio_a2v + del scale_ca_audio_a2v, shift_ca_audio_a2v + a2v_mask = perturbations.mask_like(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx, vx) + vx = vx + ( + self.audio_to_video_attn( + vx_scaled, + context=ax_scaled, + pe=video.cross_positional_embeddings, + k_pe=audio.cross_positional_embeddings, + mask=audio.key_padding_mask, + enforce_mask=audio.key_padding_mask is not None, + ) + * gate_out_a2v + * a2v_mask + ) + del gate_out_a2v, a2v_mask, vx_scaled, ax_scaled + + if run_v2a and not perturbations.all_in_batch(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx): + scale_ca_audio_v2a, shift_ca_audio_v2a, gate_out_v2a = self.get_av_ca_ada_values( + self.scale_shift_table_a2v_ca_audio, + ax.shape[0], + audio.cross_scale_shift_timestep, + audio.cross_gate_timestep, + slice(2, 4), + ) + ax_scaled = ax_norm3 * (1 + scale_ca_audio_v2a) + shift_ca_audio_v2a + del scale_ca_audio_v2a, shift_ca_audio_v2a + scale_ca_video_v2a, shift_ca_video_v2a, _ = self.get_av_ca_ada_values( + self.scale_shift_table_a2v_ca_video, + vx.shape[0], + video.cross_scale_shift_timestep, + video.cross_gate_timestep, + slice(2, 4), + ) + vx_scaled = vx_norm3 * (1 + scale_ca_video_v2a) + shift_ca_video_v2a + del scale_ca_video_v2a, shift_ca_video_v2a + v2a_mask = perturbations.mask_like(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx, ax) + ax = ax + ( + self.video_to_audio_attn( + ax_scaled, + context=vx_scaled, + pe=audio.cross_positional_embeddings, + k_pe=video.cross_positional_embeddings, + mask=video.key_padding_mask, + enforce_mask=video.key_padding_mask is not None, + ) + * gate_out_v2a + * v2a_mask + ) + del gate_out_v2a, v2a_mask, ax_scaled, vx_scaled + + del vx_norm3, ax_norm3 + + if run_vx: + vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values( + self.scale_shift_table, vx.shape[0], video.timesteps, slice(3, 6) + ) + vx_scaled = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_mlp) + vshift_mlp + vx = vx + self.ff(vx_scaled) * vgate_mlp + + del vshift_mlp, vscale_mlp, vgate_mlp, vx_scaled + + if run_ax: + ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values( + self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(3, 6) + ) + ax_scaled = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_mlp) + ashift_mlp + ax = ax + self.audio_ff(ax_scaled) * agate_mlp + + del ashift_mlp, ascale_mlp, agate_mlp, ax_scaled + + return replace(video, x=vx) if video is not None else None, replace(audio, x=ax) if audio is not None else None + + +def apply_cross_attention_adaln( + x: torch.Tensor, + context: torch.Tensor, + attn: Attention, + q_shift: torch.Tensor, + q_scale: torch.Tensor, + q_gate: torch.Tensor, + prompt_scale_shift_table: torch.Tensor, + prompt_timestep: torch.Tensor, + context_mask: torch.Tensor | None = None, + norm_eps: float = 1e-6, +) -> torch.Tensor: + batch_size = x.shape[0] + shift_kv, scale_kv = ( + prompt_scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) + + prompt_timestep.reshape(batch_size, prompt_timestep.shape[1], 2, -1) + ).unbind(dim=2) + attn_input = rms_norm(x, eps=norm_eps) * (1 + q_scale) + q_shift + encoder_hidden_states = context * (1 + scale_kv) + shift_kv + return attn(attn_input, context=encoder_hidden_states, mask=context_mask) * q_gate + + +import torch + + +class PixArtAlphaTextProjection(torch.nn.Module): + """ + Projects caption embeddings using dual linear layers. + Flow: linear_1 -> activation -> linear_2 + Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py + """ + + def __init__(self, in_features: int, hidden_size: int, out_features: int | None = None, act_fn: str = "gelu_tanh"): + super().__init__() + if out_features is None: + out_features = hidden_size + self.linear_1 = torch.nn.Linear(in_features=in_features, out_features=hidden_size, bias=True) + if act_fn == "gelu_tanh": + self.act_1 = torch.nn.GELU(approximate="tanh") + elif act_fn == "silu": + self.act_1 = torch.nn.SiLU() + else: + raise ValueError(f"Unknown activation function: {act_fn}") + self.linear_2 = torch.nn.Linear(in_features=hidden_size, out_features=out_features, bias=True) + + def forward(self, caption: torch.Tensor) -> torch.Tensor: + hidden_states = self.linear_1(caption) + hidden_states = self.act_1(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + + +def create_caption_projection(transformer_config: dict, audio: bool = False) -> PixArtAlphaTextProjection: + """Create a caption projection for the transformer (V1/19B only).""" + caption_channels = transformer_config["caption_channels"] + if audio: + inner_dim = transformer_config["audio_num_attention_heads"] * transformer_config["audio_attention_head_dim"] + else: + inner_dim = transformer_config["num_attention_heads"] * transformer_config["attention_head_dim"] + return PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim) + + +from enum import Enum + +import torch + + +class LTXModelType(Enum): + AudioVideo = "ltx av model" + VideoOnly = "ltx video only model" + AudioOnly = "ltx audio only model" + + def is_video_enabled(self) -> bool: + return self in (LTXModelType.AudioVideo, LTXModelType.VideoOnly) + + def is_audio_enabled(self) -> bool: + return self in (LTXModelType.AudioVideo, LTXModelType.AudioOnly) + + +class LTXModel(torch.nn.Module): + """ + LTX model transformer implementation. + This class implements the transformer blocks for the LTX model. + """ + + def __init__( # noqa: PLR0913 + self, + *, + model_type: LTXModelType = LTXModelType.AudioVideo, + num_attention_heads: int = 32, + attention_head_dim: int = 128, + in_channels: int = 128, + out_channels: int = 128, + num_layers: int = 48, + cross_attention_dim: int = 4096, + norm_eps: float = 1e-06, + positional_embedding_theta: float = 10000.0, + positional_embedding_max_pos: list[int] | None = None, + timestep_scale_multiplier: int = 1000, + use_middle_indices_grid: bool = True, + audio_num_attention_heads: int = 32, + audio_attention_head_dim: int = 64, + audio_in_channels: int = 128, + audio_out_channels: int = 128, + audio_cross_attention_dim: int = 2048, + audio_positional_embedding_max_pos: list[int] | None = None, + av_ca_timestep_scale_multiplier: int = 1, + rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, + double_precision_rope: bool = False, + apply_gated_attention: bool = False, + ff_bias: bool = True, + audio_ff_bias: bool = True, + use_keyframes_abs_pos_embedding: bool = False, + caption_projection: torch.nn.Module | None = None, + audio_caption_projection: torch.nn.Module | None = None, + cross_attention_adaln: bool = False, + ): + super().__init__() + self._enable_gradient_checkpointing = False + self.cross_attention_adaln = cross_attention_adaln + self.use_middle_indices_grid = use_middle_indices_grid + self.rope_type = rope_type + self.double_precision_rope = double_precision_rope + self.timestep_scale_multiplier = timestep_scale_multiplier + self.positional_embedding_theta = positional_embedding_theta + self.model_type = model_type + self.use_keyframes_abs_pos_embedding = use_keyframes_abs_pos_embedding + cross_pe_max_pos = None + if model_type.is_video_enabled(): + if positional_embedding_max_pos is None: + positional_embedding_max_pos = [20, 2048, 2048] + self.positional_embedding_max_pos = positional_embedding_max_pos + self.num_attention_heads = num_attention_heads + self.inner_dim = num_attention_heads * attention_head_dim + if use_keyframes_abs_pos_embedding: + self.keyframes_abs_pos_embedding = torch.nn.Parameter(torch.zeros(1, self.inner_dim)) + self._init_video( + in_channels=in_channels, + out_channels=out_channels, + norm_eps=norm_eps, + caption_projection=caption_projection, + ) + + if model_type.is_audio_enabled(): + if audio_positional_embedding_max_pos is None: + audio_positional_embedding_max_pos = [20] + self.audio_positional_embedding_max_pos = audio_positional_embedding_max_pos + self.audio_num_attention_heads = audio_num_attention_heads + self.audio_inner_dim = self.audio_num_attention_heads * audio_attention_head_dim + self._init_audio( + in_channels=audio_in_channels, + out_channels=audio_out_channels, + norm_eps=norm_eps, + caption_projection=audio_caption_projection, + ) + + if model_type.is_video_enabled() and model_type.is_audio_enabled(): + cross_pe_max_pos = max(self.positional_embedding_max_pos[0], self.audio_positional_embedding_max_pos[0]) + self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier + self.audio_cross_attention_dim = audio_cross_attention_dim + self._init_audio_video(num_scale_shift_values=4) + + self._init_preprocessors(cross_pe_max_pos) + # Initialize transformer blocks + self._init_transformer_blocks( + num_layers=num_layers, + attention_head_dim=attention_head_dim if model_type.is_video_enabled() else 0, + cross_attention_dim=cross_attention_dim, + audio_attention_head_dim=audio_attention_head_dim if model_type.is_audio_enabled() else 0, + audio_cross_attention_dim=audio_cross_attention_dim, + norm_eps=norm_eps, + apply_gated_attention=apply_gated_attention, + ff_bias=ff_bias, + audio_ff_bias=audio_ff_bias, + ) + self.device_mesh: DeviceMesh | None = None + self.usp_flag = False + + def enable_usp(self, device_mesh: DeviceMesh) -> None: + """Enable Ulysses sequence parallelism for self- and audio/video cross-attention.""" + strategy = get_attention_strategy(device_mesh) + if strategy != "ulysses": + raise NotImplementedError(f"LTX-2.5 supports Ulysses sequence parallelism only, got {strategy!r}") + world_size = get_ulysses_world_size(device_mesh) + for name, heads in ( + ("video", getattr(self, "num_attention_heads", 0)), + ("audio", getattr(self, "audio_num_attention_heads", 0)), + ): + if heads and heads % world_size: + raise ValueError( + f"LTX-2.5 {name} attention heads ({heads}) must be divisible by Ulysses degree ({world_size})" + ) + process_group = get_ulysses_group(device_mesh) + for block in self.transformer_blocks: + for attention_name in ("attn1", "audio_attn1", "audio_to_video_attn", "video_to_audio_attn"): + attention = getattr(block, attention_name, None) + if attention is not None: + attention.set_ulysses_group(process_group) + self.device_mesh = device_mesh + self.usp_flag = True + + def _positional_sequence_dim(self, tensor: torch.Tensor, sequence_length: int) -> int: + sequence_dim = 2 if self.rope_type == LTXRopeType.SPLIT else 1 + if tensor.ndim <= sequence_dim or tensor.shape[sequence_dim] != sequence_length: + raise ValueError( + f"Expected LTX-2.5 {self.rope_type.value} positional embeddings to have sequence length " + f"{sequence_length} at dimension {sequence_dim}, got {tuple(tensor.shape)}" + ) + return sequence_dim + + def _shard_transformer_args(self, args: TransformerArgs) -> tuple[TransformerArgs, int]: + if self.device_mesh is None: + raise RuntimeError("LTX-2.5 Ulysses is enabled without a device mesh") + sequence_length = args.x.shape[1] + world_size = get_ulysses_world_size(self.device_mesh) + padded_length = ((sequence_length + world_size - 1) // world_size) * world_size + pad_length = padded_length - sequence_length + + tensors = [args.x] + sequence_dims = [1] + for tensor in (args.timesteps, args.embedded_timestep, args.cross_scale_shift_timestep): + if tensor is not None and tensor.ndim > 1 and tensor.shape[1] == sequence_length: + tensors.append(tensor) + sequence_dims.append(1) + for embeddings in (args.positional_embeddings, args.cross_positional_embeddings): + if embeddings is not None: + for tensor in embeddings: + tensors.append(tensor) + sequence_dims.append(self._positional_sequence_dim(tensor, sequence_length)) + sequence_parallel_shard(self.device_mesh, tensors, sequence_dims) + + self_attention_mask = args.self_attention_mask + key_padding_mask = None + if pad_length: + mask_value = torch.finfo(args.x.dtype).min + key_padding_mask = torch.zeros( + args.x.shape[0], 1, 1, padded_length, device=args.x.device, dtype=args.x.dtype + ) + key_padding_mask[..., sequence_length:] = mask_value + if self_attention_mask is not None: + padded_mask = torch.zeros( + self_attention_mask.shape[0], + self_attention_mask.shape[1], + padded_length, + padded_length, + device=self_attention_mask.device, + dtype=self_attention_mask.dtype, + ) + padded_mask[..., :sequence_length, :sequence_length] = self_attention_mask + padded_mask[..., sequence_length:] = mask_value + self_attention_mask = padded_mask + else: + self_attention_mask = key_padding_mask + return replace( + args, + self_attention_mask=self_attention_mask, + key_padding_mask=key_padding_mask, + ), sequence_length + + @property + def _adaln_embedding_coefficient(self) -> int: + return adaln_embedding_coefficient(self.cross_attention_adaln) + + def _keyframes_embedding(self) -> torch.Tensor | None: + """Resolve the checkpoint-backed keyframe marker on every forward setup.""" + return getattr(self, "keyframes_abs_pos_embedding", None) + + def _init_video( + self, + in_channels: int, + out_channels: int, + norm_eps: float, + caption_projection: torch.nn.Module | None = None, + ) -> None: + """Initialize video-specific components.""" + # Video input components + self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True) + if caption_projection is not None: + self.caption_projection = caption_projection + + self.adaln_single = AdaLayerNormSingle(self.inner_dim, embedding_coefficient=self._adaln_embedding_coefficient) + + self.prompt_adaln_single = ( + AdaLayerNormSingle(self.inner_dim, embedding_coefficient=2) if self.cross_attention_adaln else None + ) + + # Video output components + self.scale_shift_table = torch.nn.Parameter(torch.empty(2, self.inner_dim)) + self.norm_out = LayerNorm(self.inner_dim, eps=norm_eps, elementwise_affine=False) + self.proj_out = torch.nn.Linear(self.inner_dim, out_channels) + + def _init_audio( + self, + in_channels: int, + out_channels: int, + norm_eps: float, + caption_projection: torch.nn.Module | None = None, + ) -> None: + """Initialize audio-specific components.""" + + # Audio input components + self.audio_patchify_proj = torch.nn.Linear(in_channels, self.audio_inner_dim, bias=True) + if caption_projection is not None: + self.audio_caption_projection = caption_projection + + self.audio_adaln_single = AdaLayerNormSingle( + self.audio_inner_dim, + embedding_coefficient=self._adaln_embedding_coefficient, + ) + + self.audio_prompt_adaln_single = ( + AdaLayerNormSingle(self.audio_inner_dim, embedding_coefficient=2) if self.cross_attention_adaln else None + ) + + # Audio output components + self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(2, self.audio_inner_dim)) + self.audio_norm_out = LayerNorm(self.audio_inner_dim, eps=norm_eps, elementwise_affine=False) + self.audio_proj_out = torch.nn.Linear(self.audio_inner_dim, out_channels) + + def _init_audio_video( + self, + num_scale_shift_values: int, + ) -> None: + """Initialize audio-video cross-attention components.""" + self.av_ca_video_scale_shift_adaln_single = AdaLayerNormSingle( + self.inner_dim, + embedding_coefficient=num_scale_shift_values, + ) + + self.av_ca_audio_scale_shift_adaln_single = AdaLayerNormSingle( + self.audio_inner_dim, + embedding_coefficient=num_scale_shift_values, + ) + + self.av_ca_a2v_gate_adaln_single = AdaLayerNormSingle( + self.inner_dim, + embedding_coefficient=1, + ) + + self.av_ca_v2a_gate_adaln_single = AdaLayerNormSingle( + self.audio_inner_dim, + embedding_coefficient=1, + ) + + def _init_preprocessors( + self, + cross_pe_max_pos: int | None = None, + ) -> None: + """Initialize preprocessors for LTX.""" + + if self.model_type.is_video_enabled() and self.model_type.is_audio_enabled(): + self.video_args_preprocessor = MultiModalTransformerArgsPreprocessor( + patchify_proj=self.patchify_proj, + adaln=self.adaln_single, + cross_scale_shift_adaln=self.av_ca_video_scale_shift_adaln_single, + cross_gate_adaln=self.av_ca_a2v_gate_adaln_single, + inner_dim=self.inner_dim, + max_pos=self.positional_embedding_max_pos, + num_attention_heads=self.num_attention_heads, + cross_pe_max_pos=cross_pe_max_pos, + use_middle_indices_grid=self.use_middle_indices_grid, + audio_cross_attention_dim=self.audio_cross_attention_dim, + timestep_scale_multiplier=self.timestep_scale_multiplier, + double_precision_rope=self.double_precision_rope, + positional_embedding_theta=self.positional_embedding_theta, + rope_type=self.rope_type, + av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier, + caption_projection=getattr(self, "caption_projection", None), + prompt_adaln=getattr(self, "prompt_adaln_single", None), + keyframes_embedding_provider=self._keyframes_embedding, + ) + self.audio_args_preprocessor = MultiModalTransformerArgsPreprocessor( + patchify_proj=self.audio_patchify_proj, + adaln=self.audio_adaln_single, + cross_scale_shift_adaln=self.av_ca_audio_scale_shift_adaln_single, + cross_gate_adaln=self.av_ca_v2a_gate_adaln_single, + inner_dim=self.audio_inner_dim, + max_pos=self.audio_positional_embedding_max_pos, + num_attention_heads=self.audio_num_attention_heads, + cross_pe_max_pos=cross_pe_max_pos, + use_middle_indices_grid=self.use_middle_indices_grid, + audio_cross_attention_dim=self.audio_cross_attention_dim, + timestep_scale_multiplier=self.timestep_scale_multiplier, + double_precision_rope=self.double_precision_rope, + positional_embedding_theta=self.positional_embedding_theta, + rope_type=self.rope_type, + av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier, + caption_projection=getattr(self, "audio_caption_projection", None), + prompt_adaln=getattr(self, "audio_prompt_adaln_single", None), + ) + elif self.model_type.is_video_enabled(): + self.video_args_preprocessor = TransformerArgsPreprocessor( + patchify_proj=self.patchify_proj, + adaln=self.adaln_single, + inner_dim=self.inner_dim, + max_pos=self.positional_embedding_max_pos, + num_attention_heads=self.num_attention_heads, + use_middle_indices_grid=self.use_middle_indices_grid, + timestep_scale_multiplier=self.timestep_scale_multiplier, + double_precision_rope=self.double_precision_rope, + positional_embedding_theta=self.positional_embedding_theta, + rope_type=self.rope_type, + caption_projection=getattr(self, "caption_projection", None), + prompt_adaln=getattr(self, "prompt_adaln_single", None), + keyframes_embedding_provider=self._keyframes_embedding, + ) + elif self.model_type.is_audio_enabled(): + self.audio_args_preprocessor = TransformerArgsPreprocessor( + patchify_proj=self.audio_patchify_proj, + adaln=self.audio_adaln_single, + inner_dim=self.audio_inner_dim, + max_pos=self.audio_positional_embedding_max_pos, + num_attention_heads=self.audio_num_attention_heads, + use_middle_indices_grid=self.use_middle_indices_grid, + timestep_scale_multiplier=self.timestep_scale_multiplier, + double_precision_rope=self.double_precision_rope, + positional_embedding_theta=self.positional_embedding_theta, + rope_type=self.rope_type, + caption_projection=getattr(self, "audio_caption_projection", None), + prompt_adaln=getattr(self, "audio_prompt_adaln_single", None), + ) + + def _init_transformer_blocks( + self, + num_layers: int, + attention_head_dim: int, + cross_attention_dim: int, + audio_attention_head_dim: int, + audio_cross_attention_dim: int, + norm_eps: float, + apply_gated_attention: bool, + ff_bias: bool, + audio_ff_bias: bool, + ) -> None: + """Initialize transformer blocks for LTX.""" + video_config = ( + TransformerConfig( + dim=self.inner_dim, + heads=self.num_attention_heads, + d_head=attention_head_dim, + context_dim=cross_attention_dim, + apply_gated_attention=apply_gated_attention, + cross_attention_adaln=self.cross_attention_adaln, + ff_bias=ff_bias, + ) + if self.model_type.is_video_enabled() + else None + ) + audio_config = ( + TransformerConfig( + dim=self.audio_inner_dim, + heads=self.audio_num_attention_heads, + d_head=audio_attention_head_dim, + context_dim=audio_cross_attention_dim, + apply_gated_attention=apply_gated_attention, + cross_attention_adaln=self.cross_attention_adaln, + ff_bias=audio_ff_bias, + ) + if self.model_type.is_audio_enabled() + else None + ) + self.transformer_blocks = torch.nn.ModuleList( + [ + BasicAVTransformerBlock( + idx=idx, + video=video_config, + audio=audio_config, + rope_type=self.rope_type, + norm_eps=norm_eps, + ) + for idx in range(num_layers) + ] + ) + + def set_gradient_checkpointing(self, enable: bool) -> None: + """Enable or disable gradient checkpointing for transformer blocks. + Gradient checkpointing trades compute for memory by recomputing activations + during the backward pass instead of storing them. This can significantly + reduce memory usage at the cost of ~20-30% slower training. + Args: + enable: Whether to enable gradient checkpointing + """ + self._enable_gradient_checkpointing = enable + + def _process_transformer_blocks( + self, + video: TransformerArgs | None, + audio: TransformerArgs | None, + perturbations: BatchedPerturbationConfig, + ) -> tuple[TransformerArgs, TransformerArgs]: + """Process transformer blocks for LTXAV.""" + + # Process transformer blocks + for block in self.transformer_blocks: + if self._enable_gradient_checkpointing and self.training: + # Use gradient checkpointing to save memory during training. + # With use_reentrant=False, we can pass dataclasses directly - + # PyTorch will track all tensor leaves in the computation graph. + video, audio = torch.utils.checkpoint.checkpoint( + block, + video, + audio, + perturbations, + use_reentrant=False, + ) + else: + video, audio = block( + video=video, + audio=audio, + perturbations=perturbations, + ) + + return video, audio + + def _process_output( + self, + scale_shift_table: torch.Tensor, + norm_out: LayerNorm, + proj_out: torch.nn.Linear, + x: torch.Tensor, + embedded_timestep: torch.Tensor, + ) -> torch.Tensor: + """Process output for LTXV.""" + # Apply scale-shift modulation + scale_shift_values = ( + scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) + embedded_timestep[:, :, None] + ) + shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] + + x = norm_out(x) + x = x * (1 + scale) + shift + x = proj_out(x) + return x + + def forward( + self, video: Modality | None, audio: Modality | None, perturbations: BatchedPerturbationConfig + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Forward pass for LTX models. + Returns: + Processed output tensors + """ + if not self.model_type.is_video_enabled() and video is not None: + raise ValueError("Video is not enabled for this model") + if not self.model_type.is_audio_enabled() and audio is not None: + raise ValueError("Audio is not enabled for this model") + + video_args = self.video_args_preprocessor.prepare(video, audio) if video is not None else None + audio_args = self.audio_args_preprocessor.prepare(audio, video) if audio is not None else None + video_sequence_length = None + audio_sequence_length = None + if self.usp_flag: + if video_args is not None: + video_args, video_sequence_length = self._shard_transformer_args(video_args) + if audio_args is not None: + audio_args, audio_sequence_length = self._shard_transformer_args(audio_args) + # Process transformer blocks + video_out, audio_out = self._process_transformer_blocks( + video=video_args, + audio=audio_args, + perturbations=perturbations, + ) + + # Process output + vx = ( + self._process_output( + self.scale_shift_table, self.norm_out, self.proj_out, video_out.x, video_out.embedded_timestep + ) + if video_out is not None + else None + ) + ax = ( + self._process_output( + self.audio_scale_shift_table, + self.audio_norm_out, + self.audio_proj_out, + audio_out.x, + audio_out.embedded_timestep, + ) + if audio_out is not None + else None + ) + if self.usp_flag: + tensors = [] + sequence_dims = [] + sequence_lengths = [] + if vx is not None and video_sequence_length is not None: + tensors.append(vx) + sequence_dims.append(1) + sequence_lengths.append(video_sequence_length) + if ax is not None and audio_sequence_length is not None: + tensors.append(ax) + sequence_dims.append(1) + sequence_lengths.append(audio_sequence_length) + unsharded = sequence_parallel_unshard(self.device_mesh, tensors, sequence_dims, sequence_lengths) + index = 0 + if vx is not None: + vx = unsharded[index] + index += 1 + if ax is not None: + ax = unsharded[index] + return vx, ax + + +class X0Model(torch.nn.Module): + """ + X0 model implementation. + Returns fully denoised outputs based on the velocities produced by the base model. + Applies scaled denoising to the video and audio according to the timesteps = sigma * denoising_mask. + """ + + def __init__(self, velocity_model: LTXModel): + super().__init__() + self.velocity_model = velocity_model + + def forward( + self, + video: Modality | None, + audio: Modality | None, + perturbations: BatchedPerturbationConfig, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """ + Denoise the video and audio according to the sigma. + Returns: + Denoised video and audio + """ + vx, ax = self.velocity_model(video, audio, perturbations) + denoised_video = to_denoised(video.latent, vx, video.timesteps) if vx is not None else None + denoised_audio = to_denoised(audio.latent, ax, audio.timesteps) if ax is not None else None + return denoised_video, denoised_audio + + +def build_ltx_model_from_config(config: dict) -> LTXModel: + """Build the registered LTX transformer from a config dict. + + This mirrors the style used by other model definitions (e.g. WanVideo) by + avoiding a dedicated `*Configurator` class and keeping the construction + logic as a simple module-level helper. + """ + caption_projection, audio_caption_projection = _build_caption_projections(config, is_av=True) + transformer_config = config.get("transformer", {}) + check_config_value(transformer_config, "dropout", 0.0) + check_config_value(transformer_config, "attention_bias", True) + check_config_value(transformer_config, "num_vector_embeds", None) + check_config_value(transformer_config, "activation_fn", "gelu-approximate") + check_config_value(transformer_config, "num_embeds_ada_norm", 1000) + check_config_value(transformer_config, "use_linear_projection", False) + check_config_value(transformer_config, "only_cross_attention", False) + check_config_value(transformer_config, "cross_attention_norm", True) + check_config_value(transformer_config, "double_self_attention", False) + check_config_value(transformer_config, "upcast_attention", False) + check_config_value(transformer_config, "standardization_norm", "rms_norm") + check_config_value(transformer_config, "norm_elementwise_affine", False) + check_config_value(transformer_config, "qk_norm", "rms_norm") + check_config_value(transformer_config, "positional_embedding_type", "rope") + check_config_value(transformer_config, "use_audio_video_cross_attention", True) + check_config_value(transformer_config, "share_ff", False) + check_config_value(transformer_config, "av_cross_ada_norm", True) + check_config_value(transformer_config, "use_middle_indices_grid", True) + return LTXModel( + model_type=LTXModelType.AudioVideo, + num_attention_heads=transformer_config.get("num_attention_heads", 32), + attention_head_dim=transformer_config.get("attention_head_dim", 128), + in_channels=transformer_config.get("in_channels", 128), + out_channels=transformer_config.get("out_channels", 128), + num_layers=transformer_config.get("num_layers", 48), + cross_attention_dim=transformer_config.get("cross_attention_dim", 4096), + norm_eps=transformer_config.get("norm_eps", 1e-06), + positional_embedding_theta=transformer_config.get("positional_embedding_theta", 10000.0), + positional_embedding_max_pos=transformer_config.get("positional_embedding_max_pos", [20, 2048, 2048]), + timestep_scale_multiplier=transformer_config.get("timestep_scale_multiplier", 1000), + use_middle_indices_grid=transformer_config.get("use_middle_indices_grid", True), + audio_num_attention_heads=transformer_config.get("audio_num_attention_heads", 32), + audio_attention_head_dim=transformer_config.get("audio_attention_head_dim", 64), + audio_in_channels=transformer_config.get("audio_in_channels", 128), + audio_out_channels=transformer_config.get("audio_out_channels", 128), + audio_cross_attention_dim=transformer_config.get("audio_cross_attention_dim", 2048), + audio_positional_embedding_max_pos=transformer_config.get("audio_positional_embedding_max_pos", [20]), + av_ca_timestep_scale_multiplier=transformer_config.get("av_ca_timestep_scale_multiplier", 1), + rope_type=LTXRopeType(transformer_config.get("rope_type", "interleaved")), + double_precision_rope=transformer_config.get("frequencies_precision", False) == "float64", + apply_gated_attention=transformer_config.get("apply_gated_attention", False), + ff_bias=transformer_config.get("ff_bias", True), + audio_ff_bias=transformer_config.get("audio_ff_bias", True), + use_keyframes_abs_pos_embedding=transformer_config.get("use_keyframes_abs_pos_embedding", False), + caption_projection=caption_projection, + audio_caption_projection=audio_caption_projection, + cross_attention_adaln=transformer_config.get("cross_attention_adaln", False), + ) + + +def _build_caption_projections(config: dict, is_av: bool) -> tuple[torch.nn.Module | None, torch.nn.Module | None]: + transformer_config = config.get("transformer", {}) + if transformer_config.get("caption_proj_before_connector", False): + return None, None + with torch.device("meta"): + caption_projection = create_caption_projection(transformer_config) + audio_caption_projection = create_caption_projection(transformer_config, audio=True) if is_av else None + return caption_projection, audio_caption_projection + + +class LTXVideoTransformer(BaseModel): + """Deprecated compatibility wrapper excluded from the LTX-2.5 public surface.""" + + def __init__(self) -> None: + super().__init__() + raise RuntimeError("LTXVideoTransformer is not available in the isolated LTX-2.5 implementation") + # Avoid registering shared modules twice (e.g. aliasing `transformer_blocks`), which would + # duplicate keys in `state_dict()` and break strict checkpoint loading. + self.layer_name_list = ["velocity_model.transformer_blocks"] + + def set_attention_config(self, attention_config: AttentionConfig) -> None: + """Set attention implementation configuration for all LTX attention blocks.""" + super().set_attention_config(attention_config) + logger.info(f"ltx dit set attention config to {attention_config.attn_impl}") + Attention.attention_config = attention_config + + @staticmethod + def state_dict_converter(): + return LTXVideoTransformerStateDictConverter() + + def get_fsdp_module_names(self) -> list[str]: + return ["velocity_model.transformer_blocks"] + + def get_tp_plan(self) -> dict: + raise NotImplementedError("Tensor parallelism plan is not implemented for LTXVideoTransformer yet.") + + def enable_quant(self, quant_type: str | torch.dtype) -> None: + from telefuser.core.config import QuantConfig, QuantType + + if isinstance(quant_type, QuantConfig) and quant_type.quant_type == QuantType.BNB_NF4: + from telefuser.ops.bnb_nf4_linear import replace_linear_layers_with_bnb_nf4 + + replaced = replace_linear_layers_with_bnb_nf4( + self.velocity_model.transformer_blocks, + compute_dtype=torch.bfloat16, + include_names=quant_type.quantize_modules, + exclude_names=quant_type.skip_modules, + ) + logger.info(f"BNB NF4 converted {replaced} Linear layers in LTX transformer blocks") + self.quant_type = quant_type.quant_type + return + if isinstance(quant_type, QuantConfig) and quant_type.quant_type == QuantType.TORCHAO_FP8: + from telefuser.ops.torchao_fp8_linear import replace_linear_layers_with_torchao_fp8 + + replaced = replace_linear_layers_with_torchao_fp8( + self.velocity_model.transformer_blocks, + include_names=quant_type.quantize_modules, + exclude_names=quant_type.skip_modules, + ) + logger.info(f"TorchAO FP8 converted {replaced} Linear layers in LTX transformer blocks") + self.quant_type = quant_type.quant_type + return + # Keep a flag for downstream pipeline logic (e.g., FSDP buffer conversion). + self.quant_type = quant_type + + def enable_usp(self) -> None: + # Interface parity with other DiT models. LTX pipeline does not currently use sequence-parallel. + self.usp_flag = True + + def compile(self, mode: str = "blocks", **kwargs) -> None: + """Compile model for better performance with torch.compile. + + Args: + mode: Compilation mode: + - "blocks": Compile only _process_transformer_blocks (default, most effective) + - "full": Compile entire forward method + **kwargs: Arguments passed to torch.compile() + """ + # Import mark_static from torch._dynamo + try: + from torch._dynamo import mark_static + + # Mark module classes as static (instance attributes won't change after compile) + mark_static(LTXVideoTransformer) + mark_static(LTXModel) + mark_static(BasicAVTransformerBlock) + mark_static(Attention) + except ImportError: + logger.warning("torch._dynamo.mark_static not available, skipping static marking") + + # Compile based on mode + if mode == "blocks": + original_fn = self.velocity_model._process_transformer_blocks + self.velocity_model._process_transformer_blocks = torch.compile(original_fn, **kwargs) + logger.info(f"LTXVideoTransformer compiled: mode={mode}") + elif mode == "full": + # Store original forward for fallback + self._original_forward = self.forward + self.forward = torch.compile(self.forward, **kwargs) + logger.info(f"LTXVideoTransformer compiled: mode={mode}") + else: + raise ValueError(f"Unknown compile mode: {mode}") + + def enable_async_offload(self, device: torch.device, offload_config: OffloadConfig) -> None: + raise NotImplementedError("Async offload is not implemented for LTXVideoTransformer yet.") + + def enable_sequential_cpu_offload( + self, + device: torch.device, + torch_dtype: torch.dtype, + max_num_param: int | None = None, + vram_limit: int | None = None, + ) -> None: + """Enable sequential CPU offloading for memory efficiency.""" + + from telefuser.offload import AutoWrappedLinear, AutoWrappedModule, enable_sequential_cpu_offload + + dtype = next(iter(self.parameters())).dtype + enable_sequential_cpu_offload( + self, + module_map={ + torch.nn.Linear: AutoWrappedLinear, + torch.nn.LayerNorm: AutoWrappedModule, + }, + module_config=dict( + offload_dtype=dtype, + offload_device="cpu", + onload_dtype=dtype, + onload_device=device, + computation_dtype=torch_dtype, + computation_device=device, + ), + max_num_param=max_num_param, + overflow_module_config=dict( + offload_dtype=dtype, + offload_device="cpu", + onload_dtype=dtype, + onload_device="cpu", + computation_dtype=torch_dtype, + computation_device=device, + ), + vram_limit=vram_limit, + ) + + def forward( + self, + video: Modality | None, + audio: Modality | None, + perturbations: BatchedPerturbationConfig, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + vx, ax = self.velocity_model(video, audio, perturbations) + denoised_video = to_denoised(video.latent, vx, video.timesteps) if vx is not None else None + denoised_audio = to_denoised(audio.latent, ax, audio.timesteps) if ax is not None else None + return denoised_video, denoised_audio + + +class LTXVideoTransformerStateDictConverter: + """Convert the shared LTX checkpoint into the registered transformer layout.""" + + _MODEL_PREFIX = "model.diffusion_model." + _EMBEDDINGS_CONNECTOR_PREFIXES = ( + # Keep both singular/plural variants for robustness across checkpoint exports. + "model.diffusion_model.video_embeddings_connector", + "model.diffusion_model.audio_embeddings_connector", + "model.diffusion_model.video_embedding_connector", + "model.diffusion_model.audio_embedding_connector", + ) + + def from_official(self, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + converted_state_dict: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if key.startswith(self._EMBEDDINGS_CONNECTOR_PREFIXES): + continue + if key.startswith(self._MODEL_PREFIX): + converted_state_dict[f"velocity_model.{key.removeprefix(self._MODEL_PREFIX)}"] = value + return converted_state_dict + + def from_diffusers(self, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + raise NotImplementedError("LTXVideoTransformer only supports the civitai-style single-file checkpoint.") + + +_LTX25_TRANSFORMER_PREFIX = "model.diffusion_model." +_LTX25_EMBEDDING_PREFIXES = ( + "video_embeddings_connector.", + "audio_embeddings_connector.", + "video_aggregate_embed.", + "audio_aggregate_embed.", +) + + +def build_ltx25_av_model(config: dict[str, Any]) -> LTXModel: + """Build the LTX-2.5 AV denoiser strictly from checkpoint metadata.""" + transformer_config = config.get("transformer") + if not isinstance(transformer_config, dict): + raise ValueError("LTX-2.5 transformer checkpoint is missing object transformer config") + required = { + "num_layers": 48, + "rope_type": "split", + "apply_gated_attention": True, + "ff_bias": False, + "caption_proj_before_connector": True, + } + mismatches = { + key: (transformer_config.get(key), expected) + for key, expected in required.items() + if transformer_config.get(key) != expected + } + if mismatches: + raise ValueError(f"Unsupported LTX-2.5 AV transformer config: {mismatches}") + return build_ltx_model_from_config(config) + + +def ltx25_transformer_key_to_model_key(key: str) -> str | None: + """Map a split LTX-2.5 checkpoint key into the isolated AV denoiser state dict.""" + if not key.startswith(_LTX25_TRANSFORMER_PREFIX): + return None + suffix = key.removeprefix(_LTX25_TRANSFORMER_PREFIX) + if suffix.startswith(_LTX25_EMBEDDING_PREFIXES): + return None + return f"velocity_model.{suffix}" + + +def ltx25_transformer_checkpoint_key_coverage( + checkpoint_path: str | Path, + model_keys: set[str], +) -> tuple[set[str], set[str]]: + """Return unexplained split-checkpoint keys and missing isolated-model keys.""" + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as checkpoint: + mapped = { + target for key in checkpoint.keys() if (target := ltx25_transformer_key_to_model_key(key)) is not None + } + return mapped - model_keys, model_keys - mapped + + +class LTX25AVTransformer(BaseModel): + """Isolated LTX-2.5 48-layer gated audio-video x0 denoiser.""" + + def __init__(self, config: dict[str, Any]) -> None: + super().__init__() + self.velocity_model = build_ltx25_av_model(config) + self.layer_name_list = ["velocity_model.transformer_blocks"] + + @classmethod + def from_checkpoint( + cls, + checkpoint_path: str | Path, + *, + device: torch.device | str = "cpu", + torch_dtype: torch.dtype = torch.bfloat16, + ) -> "LTX25AVTransformer": + """Construct and strictly load the split LTX-2.5 denoiser checkpoint.""" + from .checkpoint import inspect_checkpoint + + metadata = inspect_checkpoint(checkpoint_path) + with torch.device("meta"): + model = cls(metadata.config) + unexpected, missing = ltx25_transformer_checkpoint_key_coverage(checkpoint_path, set(model.state_dict())) + if unexpected or missing: + raise ValueError( + "LTX-2.5 transformer checkpoint coverage mismatch: " + f"unexpected={sorted(unexpected)[:5]}, missing={sorted(missing)[:5]}" + ) + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as checkpoint: + state_dict = { + target: checkpoint.get_tensor(key) + for key in checkpoint.keys() + if (target := ltx25_transformer_key_to_model_key(key)) is not None + } + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=True, assign=True) + if missing_keys or unexpected_keys: + raise ValueError( + f"LTX-2.5 transformer load mismatch: missing={missing_keys[:5]}, unexpected={unexpected_keys[:5]}" + ) + # On the supported H100 runtime, upstream selects FlashAttention 4. + # Keep this fidelity baseline explicit; the public attention dispatch + # retains its normal fallback behavior when FA4 is unavailable. + model.set_attention_config(AttentionConfig.dense_attention(AttnImplType.FLASH_ATTN_4)) + return model.to(device=device, dtype=torch_dtype).eval() + + def set_attention_config(self, attention_config: AttentionConfig) -> None: + """Configure the public TeleFuser attention dispatch for all isolated blocks.""" + super().set_attention_config(attention_config) + Attention.attention_config = attention_config + + def enable_usp(self, device_mesh: DeviceMesh) -> None: + """Enable Ulysses sequence parallelism in the isolated AV transformer.""" + self.device_mesh = device_mesh + self.velocity_model.enable_usp(device_mesh) + self.usp_flag = True + + def get_fsdp_module_names(self) -> list[str]: + """Return FSDP wrap targets relative to ``velocity_model``.""" + return ["transformer_blocks"] + + def forward( + self, + video: Modality | None, + audio: Modality | None, + perturbations: BatchedPerturbationConfig, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """Return x0 predictions for the supplied LTX-2.5 modality states.""" + velocity_video, velocity_audio = self.velocity_model(video, audio, perturbations) + denoised_video = ( + to_denoised(video.latent, velocity_video, video.timesteps) if velocity_video is not None else None + ) + denoised_audio = ( + to_denoised(audio.latent, velocity_audio, audio.timesteps) if velocity_audio is not None else None + ) + return denoised_video, denoised_audio + + +__all__ = [ + "BatchedPerturbationConfig", + "LTX25AVTransformer", + "LTXModel", + "Modality", + "Perturbation", + "PerturbationConfig", + "PerturbationType", + "X0Model", + "build_ltx25_av_model", + "ltx25_transformer_checkpoint_key_coverage", + "ltx25_transformer_key_to_model_key", +] diff --git a/telefuser/models/ltx25/video_encoder.py b/telefuser/models/ltx25/video_encoder.py new file mode 100644 index 0000000..5ed9754 --- /dev/null +++ b/telefuser/models/ltx25/video_encoder.py @@ -0,0 +1,97 @@ +"""Isolated LTX-2.5 video encoder split-checkpoint loader.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import torch +from safetensors import safe_open + +from .checkpoint import inspect_checkpoint +from .conv_video_vae import LogVarianceType, NormLayerType, PaddingModeType, VideoEncoder + + +class LTX25VideoEncoder(VideoEncoder): + """LTX-2.5 video encoder built from the split VAE metadata.""" + + @classmethod + def from_checkpoint( + cls, + checkpoint_path: str | Path, + *, + device: torch.device | str = "cpu", + torch_dtype: torch.dtype = torch.bfloat16, + ) -> "LTX25VideoEncoder": + """Construct and strictly load the video encoder from either official VAE variant.""" + checkpoint = inspect_checkpoint(checkpoint_path) + kwargs = _video_encoder_kwargs(checkpoint.config) + with torch.device("meta"): + model = cls(**kwargs) + unexpected, missing = ltx25_video_encoder_checkpoint_key_coverage(checkpoint.path, set(model.state_dict())) + if unexpected or missing: + raise ValueError( + "LTX-2.5 video encoder checkpoint coverage mismatch: " + f"unexpected={sorted(unexpected)[:5]}, missing={sorted(missing)[:5]}" + ) + state_dict = _load_video_encoder_state_dict(checkpoint.path) + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=True, assign=True) + if missing_keys or unexpected_keys: + raise ValueError( + f"LTX-2.5 video encoder load mismatch: missing={missing_keys[:5]}, unexpected={unexpected_keys[:5]}" + ) + return model.to(device=device, dtype=torch_dtype).eval() + + +def _video_encoder_kwargs(config: dict[str, Any]) -> dict[str, Any]: + vae_config = config.get("vae") + if not isinstance(vae_config, dict): + raise ValueError("LTX-2.5 VAE checkpoint is missing object vae config") + encoder_config = vae_config.get("encoder", vae_config) + if not isinstance(encoder_config, dict): + raise ValueError("LTX-2.5 VAE encoder config must be an object") + blocks = encoder_config.get("blocks", encoder_config.get("encoder_blocks")) + if not isinstance(blocks, list): + raise ValueError("LTX-2.5 VAE encoder config is missing blocks") + return { + "convolution_dimensions": encoder_config.get("dims", vae_config.get("dims", 3)), + "in_channels": encoder_config.get("in_channels", 3), + "out_channels": encoder_config.get("latent_channels", vae_config.get("latent_channels", 128)), + "encoder_blocks": blocks, + "patch_size": encoder_config.get("patch_size", 4), + "norm_layer": NormLayerType(encoder_config.get("norm_layer", "pixel_norm")), + "latent_log_var": LogVarianceType(encoder_config.get("latent_log_var", "uniform")), + "encoder_spatial_padding_mode": PaddingModeType(encoder_config.get("spatial_padding_mode", "zeros")), + } + + +def _video_encoder_target_key(key: str) -> str | None: + if key.startswith("encoder."): + return key + if key.startswith("per_channel_statistics."): + return "per_channel_statistics." + key.removeprefix("per_channel_statistics.") + return None + + +def _load_video_encoder_state_dict(checkpoint_path: str | Path) -> dict[str, torch.Tensor]: + state_dict: dict[str, torch.Tensor] = {} + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as checkpoint: + for key in checkpoint.keys(): + target = _video_encoder_target_key(key) + if target is not None: + state_dict[target.removeprefix("encoder.")] = checkpoint.get_tensor(key) + return state_dict + + +def ltx25_video_encoder_checkpoint_key_coverage( + checkpoint_path: str | Path, + model_keys: set[str], +) -> tuple[set[str], set[str]]: + """Return unexplained source keys and missing isolated video encoder keys.""" + mapped: set[str] = set() + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as checkpoint: + for key in checkpoint.keys(): + target = _video_encoder_target_key(key) + if target is not None: + mapped.add(target.removeprefix("encoder.")) + return mapped - model_keys, model_keys - mapped diff --git a/telefuser/ops/__init__.py b/telefuser/ops/__init__.py index 9e5da26..74a0659 100644 --- a/telefuser/ops/__init__.py +++ b/telefuser/ops/__init__.py @@ -14,6 +14,11 @@ from .base import CustomOp, CustomOpFunction from .custom_op import TritonKernelWrapper, register_custom_op from .moe import grouped_expert_forward, route_topk +from .neighborhood_attention import ( + configure_neighborhood_attention_kv_parallelism, + natten_available, + neighborhood_attention_3d, +) from .normalization import ( AdaLayerNormContinuous, LayerNorm, @@ -38,10 +43,14 @@ "AdaLayerNormContinuous", "fused_scale_shift", "modulate", + "configure_neighborhood_attention_kv_parallelism", "indexed_gate", "indexed_scale_shift", "route_topk", "grouped_expert_forward", + # Neighborhood attention + "natten_available", + "neighborhood_attention_3d", # Rotary "apply_rotary_emb", "apply_qk_norm_rope_neox", diff --git a/telefuser/ops/neighborhood_attention.py b/telefuser/ops/neighborhood_attention.py new file mode 100644 index 0000000..a50545f --- /dev/null +++ b/telefuser/ops/neighborhood_attention.py @@ -0,0 +1,53 @@ +"""Public dispatch for optional 3D neighborhood-attention backends.""" + +from __future__ import annotations + +import torch + +try: + import natten + + _NATTEN_AVAILABLE = True +except ImportError: # pragma: no cover + natten = None # type: ignore[assignment] + _NATTEN_AVAILABLE = False + + +def natten_available() -> bool: + """Return whether NATTEN's required CUDA extension is usable.""" + + return _NATTEN_AVAILABLE and bool(getattr(natten, "HAS_LIBNATTEN", True)) + + +def configure_neighborhood_attention_kv_parallelism(enabled: bool) -> bool: + """Configure NATTEN fused-attention KV parallelism when its CUDA extension is present.""" + if not natten_available(): + return False + configure = getattr(natten, "use_kv_parallelism_in_fused_na", None) + if configure is None: + return False + configure(enabled) + return True + + +def neighborhood_attention_3d( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + kernel_size: tuple[int, int, int], + backend: str | None = None, +) -> torch.Tensor: + """Run NATTEN 3D neighborhood attention with the framework's tensor contract. + + The model pre-scales query, so the NATTEN scale remains one. NATTEN requires + matching Q/K/V dtypes; the cast mirrors the existing flash-attention paths. + """ + if not natten_available(): + raise ImportError( + "natten is required for 3D neighborhood attention. Install a NATTEN build that includes libnatten." + ) + if query.dtype != value.dtype or key.dtype != value.dtype: + query = query.to(dtype=value.dtype) + key = key.to(dtype=value.dtype) + return natten.na3d(query, key, value, kernel_size=kernel_size, scale=1.0, backend=backend) diff --git a/telefuser/pipelines/ltx25_distilled/__init__.py b/telefuser/pipelines/ltx25_distilled/__init__.py new file mode 100644 index 0000000..3712bad --- /dev/null +++ b/telefuser/pipelines/ltx25_distilled/__init__.py @@ -0,0 +1,33 @@ +"""LTX-2.5 distilled pipeline package.""" + +from .loader import load_ltx25_distilled_modules +from .pipeline import ( + LTX25DistilledConfig, + LTX25DistilledOutput, + LTX25DistilledPipeline, + LTX25ImageCondition, + build_ltx25_distilled_config, +) +from .reference import ( + LTX25DistilledReference, + LTX25ReferenceComponents, + LTX25ReferenceImageCondition, + LTX25ReferenceRequest, + LTX25ReferenceResult, + LTX25ReferenceTrace, +) + +__all__ = [ + "LTX25DistilledReference", + "LTX25ReferenceComponents", + "LTX25ReferenceImageCondition", + "LTX25ReferenceRequest", + "LTX25ReferenceResult", + "LTX25ReferenceTrace", + "LTX25DistilledConfig", + "LTX25DistilledOutput", + "LTX25DistilledPipeline", + "LTX25ImageCondition", + "build_ltx25_distilled_config", + "load_ltx25_distilled_modules", +] diff --git a/telefuser/pipelines/ltx25_distilled/audio_decoding.py b/telefuser/pipelines/ltx25_distilled/audio_decoding.py new file mode 100644 index 0000000..8ce6210 --- /dev/null +++ b/telefuser/pipelines/ltx25_distilled/audio_decoding.py @@ -0,0 +1,25 @@ +"""Audio decoding and vocoding for LTX-2.5.""" + +from __future__ import annotations + +import torch + +from telefuser.core.base_stage import BaseStage, with_model_offload +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager + + +class LTX25AudioDecodingStage(BaseStage): + def __init__(self, module_manager: ModuleManager, config: ModelRuntimeConfig) -> None: + super().__init__("ltx25_audio_decoding", config) + self.audio_decoder: torch.nn.Module = module_manager.fetch_module("ltx25_audio_decoder") + self.vocoder: torch.nn.Module = module_manager.fetch_module("ltx25_vocoder") + self.model_names = ["audio_decoder", "vocoder"] + + @with_model_offload(["audio_decoder", "vocoder"]) + @torch.inference_mode() + def decode(self, latent: torch.Tensor) -> torch.Tensor: + return self.vocoder(self.audio_decoder(latent)).squeeze(0).float() + + +__all__ = ["LTX25AudioDecodingStage"] diff --git a/telefuser/pipelines/ltx25_distilled/core.py b/telefuser/pipelines/ltx25_distilled/core.py new file mode 100644 index 0000000..ac59af7 --- /dev/null +++ b/telefuser/pipelines/ltx25_distilled/core.py @@ -0,0 +1,183 @@ +"""Faithful LTX-2.5 distilled denoising helpers. + +This module intentionally owns only the single-GPU sampling mechanics. Stage +assembly, prompting, and decoding remain separate so the sampler can be tested +without loading the 22B transformer. +""" + +from __future__ import annotations + +from dataclasses import replace + +import torch + +from telefuser.models.ltx25.sampler import LTX25EulerAncestralStep +from telefuser.models.ltx25.transformer import LTX25AVTransformer, Modality + +from .latent import LatentState + + +def timesteps_from_mask(denoise_mask: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor: + """Apply a batch of scalar sigmas to the corresponding token masks.""" + if sigma.ndim != 1: + raise ValueError(f"sigma must have shape (batch,), got {tuple(sigma.shape)}") + return denoise_mask * sigma.view(-1, *([1] * (denoise_mask.ndim - 1))) + + +def post_process_latent(denoised: torch.Tensor, state: LatentState) -> torch.Tensor: + """Keep conditioning tokens fixed after a denoiser prediction.""" + return (denoised * state.denoise_mask + state.clean_latent.float() * (1 - state.denoise_mask)).to(denoised.dtype) + + +def modality_from_latent_state(state: LatentState, context: torch.Tensor, sigma: torch.Tensor) -> Modality: + """Translate an isolated latent state into the transformer's public input.""" + return Modality( + latent=state.latent, + sigma=sigma, + timesteps=timesteps_from_mask(state.denoise_mask, sigma), + positions=state.positions, + context=context, + context_mask=None, + attention_mask=state.attention_mask, + keyframes_mask=state.keyframes_mask, + ) + + +class LTX25SimpleDenoiser: + """One conditioned transformer call without CFG or perturbations.""" + + def __init__(self, video_context: torch.Tensor | None, audio_context: torch.Tensor | None) -> None: + self.video_context = video_context + self.audio_context = audio_context + + def __call__( + self, + transformer: LTX25AVTransformer, + video_state: LatentState | None, + audio_state: LatentState | None, + sigmas: torch.Tensor, + step_index: int, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + if video_state is None and audio_state is None: + raise ValueError("At least one latent modality must be provided") + if video_state is not None and self.video_context is None: + raise ValueError("video_context is required when video_state is present") + if audio_state is not None and self.audio_context is None: + raise ValueError("audio_context is required when audio_state is present") + + sigma = sigmas[step_index] + video = ( + modality_from_latent_state(video_state, self.video_context, sigma.expand(video_state.latent.shape[0])) + if video_state is not None + else None + ) + audio = ( + modality_from_latent_state(audio_state, self.audio_context, sigma.expand(audio_state.latent.shape[0])) + if audio_state is not None + else None + ) + return transformer(video=video, audio=audio, perturbations=None) + + +def euler_denoising_loop( + sigmas: torch.Tensor, + video_state: LatentState | None, + audio_state: LatentState | None, + transformer: LTX25AVTransformer, + denoiser: LTX25SimpleDenoiser, + *, + model_dtype: torch.dtype, +) -> tuple[LatentState | None, LatentState | None]: + """Run the deterministic rectified-flow Euler loop used by distilled stage 2.""" + for step_index in range(sigmas.numel() - 1): + denoised_video, denoised_audio = denoiser(transformer, video_state, audio_state, sigmas, step_index) + video_state = _euler_step(video_state, denoised_video, sigmas, step_index, model_dtype) + audio_state = _euler_step(audio_state, denoised_audio, sigmas, step_index, model_dtype) + return video_state, audio_state + + +def euler_ancestral_denoising_loop( + sigmas: torch.Tensor, + video_state: LatentState | None, + audio_state: LatentState | None, + transformer: LTX25AVTransformer, + denoiser: LTX25SimpleDenoiser, + *, + noise_seed: int, + stepper: LTX25EulerAncestralStep | None = None, + model_dtype: torch.dtype, +) -> tuple[LatentState | None, LatentState | None]: + """Run stage-1 ancestral Euler sampling with upstream noise ordering.""" + if video_state is None and audio_state is None: + raise ValueError("At least one latent modality must be provided") + stepper = stepper or LTX25EulerAncestralStep() + present_state = video_state if video_state is not None else audio_state + generator = torch.Generator(device=present_state.latent.device).manual_seed(noise_seed) + + for step_index in range(sigmas.numel() - 1): + denoised_video, denoised_audio = denoiser(transformer, video_state, audio_state, sigmas, step_index) + terminal_step = bool(sigmas[step_index + 1] == 0) + video_state = _ancestral_step( + video_state, denoised_video, sigmas, step_index, terminal_step, stepper, generator, model_dtype + ) + audio_state = _ancestral_step( + audio_state, denoised_audio, sigmas, step_index, terminal_step, stepper, generator, model_dtype + ) + return video_state, audio_state + + +def _euler_step( + state: LatentState | None, + denoised: torch.Tensor | None, + sigmas: torch.Tensor, + step_index: int, + model_dtype: torch.dtype, +) -> LatentState | None: + if state is None or denoised is None: + return state + denoised = post_process_latent(denoised, state) + sigma = sigmas[step_index] + sigma_next = sigmas[step_index + 1] + if bool(sigma_next == 0): + return replace(state, latent=denoised.to(model_dtype)) + # Preserve EulerDiffusionStep's BF16 velocity rounding point. The ratio + # form is algebraically equivalent but diverges from the upstream trajectory. + velocity = ((state.latent.float() - denoised.float()) / sigma.to(torch.float32).item()).to(state.latent.dtype) + latent = (state.latent.float() + velocity.float() * (sigma_next - sigma)).to(model_dtype) + return replace(state, latent=latent) + + +def _ancestral_step( + state: LatentState | None, + denoised: torch.Tensor | None, + sigmas: torch.Tensor, + step_index: int, + terminal_step: bool, + stepper: LTX25EulerAncestralStep, + generator: torch.Generator, + model_dtype: torch.dtype, +) -> LatentState | None: + if state is None or denoised is None: + return state + denoised = post_process_latent(denoised, state) + if terminal_step: + return replace(state, latent=denoised.to(model_dtype)) + noise = torch.randn(state.latent.shape, generator=generator, dtype=state.latent.dtype, device=state.latent.device) + latent = stepper.step( + sample=state.latent.float(), + denoised_sample=denoised, + sigmas=sigmas, + step_index=step_index, + noise=noise, + ) + return replace(state, latent=post_process_latent(latent, state).to(model_dtype)) + + +__all__ = [ + "LTX25SimpleDenoiser", + "euler_ancestral_denoising_loop", + "euler_denoising_loop", + "modality_from_latent_state", + "post_process_latent", + "timesteps_from_mask", +] diff --git a/telefuser/pipelines/ltx25_distilled/denoising.py b/telefuser/pipelines/ltx25_distilled/denoising.py new file mode 100644 index 0000000..7f97f27 --- /dev/null +++ b/telefuser/pipelines/ltx25_distilled/denoising.py @@ -0,0 +1,158 @@ +"""Two-phase distilled denoising for LTX-2.5.""" + +from __future__ import annotations + +import torch + +from telefuser.core.base_stage import BaseStage +from telefuser.core.config import ModelRuntimeConfig, WeightOffloadType +from telefuser.core.module_manager import ModuleManager +from telefuser.distributed import create_device_mesh_from_config +from telefuser.distributed.fsdp import shard_model_fsdp2_inference +from telefuser.models.ltx25 import LTX25AVTransformer +from telefuser.models.ltx25.sampler import LTX25_STAGE1_DISTILLED_SIGMAS, LTX25_STAGE2_DISTILLED_SIGMAS +from telefuser.offload import AsyncOffloadManager +from telefuser.platforms import current_platform +from telefuser.utils.logging import logger + +from .core import LTX25SimpleDenoiser, euler_ancestral_denoising_loop, euler_denoising_loop +from .latent import LatentState + +LatentStateInput = LatentState | dict[str, torch.Tensor | None] + + +def _restore_latent_state(state: LatentStateInput) -> LatentState: + if isinstance(state, LatentState): + return state + return LatentState( + latent=state["latent"], + denoise_mask=state["denoise_mask"], + positions=state["positions"], + clean_latent=state["clean_latent"], + attention_mask=state.get("attention_mask"), + keyframes_mask=state.get("keyframes_mask"), + ) + + +class LTX25DenoisingStage(BaseStage): + def __init__(self, module_manager: ModuleManager, config: ModelRuntimeConfig) -> None: + super().__init__("ltx25_denoising", config) + self.transformer: LTX25AVTransformer = module_manager.fetch_module("ltx25_transformer") + self.transformer.set_attention_config(config.attention_config) + self.model_names = ["transformer"] + self.empty_cache_after_call = False + self.offload_manager: AsyncOffloadManager | None = None + if config.offload_config.offload_type == WeightOffloadType.ASYNC_CPU_OFFLOAD: + self.offload_manager = AsyncOffloadManager( + self.transformer.velocity_model.transformer_blocks, + device=self.device, + pin_cpu_memory=config.offload_config.pin_cpu_memory, + offload_ratio=config.offload_config.offload_ratio, + prefetch_size=config.offload_config.prefetch_size, + ) + self.transformer.to(device=self.device, dtype=self.torch_dtype) + self.onload_models_flag = True + + def parallel_models(self) -> None: + """Configure Ulysses SP and optional block-level FSDP2 for denoising.""" + parallel_config = self.model_runtime_config.parallel_config + unsupported = { + "dp_degree": parallel_config.dp_degree, + "cfg_degree": parallel_config.cfg_degree, + "sp_ring_degree": parallel_config.sp_ring_degree, + "pp_degree": parallel_config.pp_degree, + "tp_degree": parallel_config.tp_degree, + } + invalid = {name: degree for name, degree in unsupported.items() if degree != 1} + if invalid: + raise NotImplementedError(f"LTX-2.5 does not support these parallel degrees: {invalid}") + device_mesh = create_device_mesh_from_config(parallel_config, device_type=self.device.type) + self.transformer.set_attention_config(self.model_runtime_config.attention_config) + if parallel_config.sp_ulysses_degree > 1: + self.transformer.enable_usp(device_mesh) + logger.info(f"enabled LTX-2.5 Ulysses SP degree={parallel_config.sp_ulysses_degree}") + if parallel_config.enable_fsdp: + if self.model_runtime_config.offload_config.offload_type != WeightOffloadType.NO_CPU_OFFLOAD: + raise ValueError("LTX-2.5 FSDP inference cannot be combined with model CPU offload") + logger.info(f"enabled LTX-2.5 block FSDP2 for {self.name}") + self.transformer.velocity_model = shard_model_fsdp2_inference( + module=self.transformer.velocity_model, + device_mesh=device_mesh, + wrap_module_names=self.transformer.get_fsdp_module_names(), + ) + self.onload_models_flag = True + current_platform.empty_cache() + + def _onload(self) -> None: + if not self.onload_models_flag: + self.transformer.to(self.device) + self.onload_models_flag = True + + def _offload_between_phases(self) -> None: + if self.model_runtime_config.offload_config.offload_type == WeightOffloadType.MODEL_CPU_OFFLOAD: + self.transformer.cpu() + self.onload_models_flag = False + + @torch.inference_mode() + def denoise_stage1( + self, + video: LatentStateInput, + audio: LatentStateInput, + video_context: torch.Tensor, + audio_context: torch.Tensor, + *, + noise_seed: int, + ) -> tuple[LatentState, LatentState]: + self._onload() + video = _restore_latent_state(video) + audio = _restore_latent_state(audio) + result = euler_ancestral_denoising_loop( + torch.tensor(LTX25_STAGE1_DISTILLED_SIGMAS, device=self.device), + video, + audio, + self.transformer, + LTX25SimpleDenoiser(video_context, audio_context), + noise_seed=noise_seed, + model_dtype=self.torch_dtype, + ) + self._offload_between_phases() + if result[0] is None or result[1] is None: + raise RuntimeError("LTX-2.5 stage-one denoising requires video and audio outputs") + if self.model_runtime_config.parallel_config.world_size > 1: + return result[0].clone(), result[1].clone() + return result[0], result[1] + + @torch.inference_mode() + def denoise_stage2( + self, + video: LatentStateInput, + audio: LatentStateInput, + video_context: torch.Tensor, + audio_context: torch.Tensor, + ) -> tuple[LatentState, LatentState]: + self._onload() + video = _restore_latent_state(video) + audio = _restore_latent_state(audio) + result = euler_denoising_loop( + torch.tensor(LTX25_STAGE2_DISTILLED_SIGMAS, device=self.device), + video, + audio, + self.transformer, + LTX25SimpleDenoiser(video_context, audio_context), + model_dtype=self.torch_dtype, + ) + if result[0] is None or result[1] is None: + raise RuntimeError("LTX-2.5 stage-two denoising requires video and audio outputs") + if self.model_runtime_config.parallel_config.world_size > 1: + return result[0].clone(), result[1].clone() + return result[0], result[1] + + def finish_request(self) -> None: + if self.offload_manager is not None: + self.offload_manager.release_all() + elif self.model_runtime_config.offload_config.offload_type != WeightOffloadType.NO_CPU_OFFLOAD: + self.transformer.cpu() + self.onload_models_flag = False + + +__all__ = ["LTX25DenoisingStage"] diff --git a/telefuser/pipelines/ltx25_distilled/image.py b/telefuser/pipelines/ltx25_distilled/image.py new file mode 100644 index 0000000..13f82ae --- /dev/null +++ b/telefuser/pipelines/ltx25_distilled/image.py @@ -0,0 +1,66 @@ +"""Faithful LTX-2.5 image-conditioning preprocessing.""" + +from __future__ import annotations + +from io import BytesIO +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + +try: + import av +except ImportError: # pragma: no cover - exercised by the runtime dependency check. + av = None + +from telefuser.models.ltx25.checkpoint import inspect_checkpoint + + +def default_image_crf(video_encoder_path: str | Path | None) -> int: + """Return the upstream default SDR image-conditioning CRF.""" + if video_encoder_path is None: + return 18 + return 18 if inspect_checkpoint(video_encoder_path).model_version >= (2, 4) else 33 + + +def preprocess_ltx25_image( + image: Image.Image, + height: int, + width: int, + crf: int, + *, + device: torch.device | str = "cpu", + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Match upstream SDR CRF, crop, resize, and VAE-normalization ordering.""" + pixels = np.array(image.convert("RGB"), dtype=np.uint8, copy=True) + if crf and min(pixels.shape[:2]) >= 2: + if av is None: + raise RuntimeError("PyAV is required for non-zero LTX-2.5 image conditioning CRF") + with BytesIO() as encoded: + container = av.open(encoded, "w", format="mp4") + try: + stream = container.add_stream("libx264", rate=1, options={"crf": str(crf), "preset": "veryfast"}) + encoded_height = pixels.shape[0] // 2 * 2 + encoded_width = pixels.shape[1] // 2 * 2 + stream.height, stream.width = encoded_height, encoded_width + frame = av.VideoFrame.from_ndarray(pixels[:encoded_height, :encoded_width], format="rgb24") + container.mux(stream.encode(frame.reformat(format="yuv420p"))) + container.mux(stream.encode()) + finally: + container.close() + with av.open(BytesIO(encoded.getvalue())) as decoded: + pixels = next(decoded.decode(video=0)).to_ndarray(format="rgb24") + # Upstream moves pixels before interpolation. CPU interpolation rounds differently. + tensor = torch.from_numpy(pixels).permute(2, 0, 1).unsqueeze(0).to(device=device, dtype=torch.float32) + source_height, source_width = tensor.shape[-2:] + scale = max(height / source_height, width / source_width) + resized_height, resized_width = int(np.ceil(source_height * scale)), int(np.ceil(source_width * scale)) + tensor = F.interpolate(tensor, size=(resized_height, resized_width), mode="bilinear", align_corners=False) + top, left = (resized_height - height) // 2, (resized_width - width) // 2 + return (tensor[:, :, top : top + height, left : left + width] / 127.5 - 1.0).to(dtype=dtype).unsqueeze(2) + + +__all__ = ["default_image_crf", "preprocess_ltx25_image"] diff --git a/telefuser/pipelines/ltx25_distilled/latent.py b/telefuser/pipelines/ltx25_distilled/latent.py new file mode 100644 index 0000000..7c41842 --- /dev/null +++ b/telefuser/pipelines/ltx25_distilled/latent.py @@ -0,0 +1,545 @@ +from __future__ import annotations + +import math + +"""Isolated LTX-2.5 latent patchification and conditioning helpers.""" + +from dataclasses import dataclass, field, replace +from typing import Any, Callable, NamedTuple, Protocol + +import einops +import torch +from torch._prims_common import DeviceLikeType + +from telefuser.models.ltx25.diff_vae.types import VIDEO_SCALE_FACTORS, SpatioTemporalScaleFactors, VideoLatentShape + +STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875, 0.0] +VIDEO_LATENT_CHANNELS = 128 + + +class AudioLatentShape(NamedTuple): + batch: int + channels: int + frames: int + mel_bins: int + + def to_torch_shape(self) -> torch.Size: + return torch.Size([self.batch, self.channels, self.frames, self.mel_bins]) + + def token_count(self) -> int: + return self.frames + + def mask_shape(self) -> AudioLatentShape: + return AudioLatentShape(self.batch, 1, self.frames, 1) + + @staticmethod + def from_duration( + batch: int, + duration: float, + channels: int = 8, + mel_bins: int = 16, + sample_rate: int = 16000, + hop_length: int = 160, + audio_latent_downsample_factor: int = 4, + ) -> AudioLatentShape: + latents_per_second = float(sample_rate) / float(hop_length) / float(audio_latent_downsample_factor) + return AudioLatentShape(batch, channels, round(duration * latents_per_second), mel_bins) + + +@dataclass(frozen=True) +class LatentState: + latent: torch.Tensor + denoise_mask: torch.Tensor + positions: torch.Tensor + clean_latent: torch.Tensor + attention_mask: torch.Tensor | None = None + keyframes_mask: torch.Tensor | None = None + + def clone(self) -> LatentState: + return LatentState( + latent=self.latent.clone(), + denoise_mask=self.denoise_mask.clone(), + positions=self.positions.clone(), + clean_latent=self.clean_latent.clone(), + attention_mask=self.attention_mask.clone() if self.attention_mask is not None else None, + keyframes_mask=self.keyframes_mask.clone() if self.keyframes_mask is not None else None, + ) + + +class Patchifier(Protocol): + def patchify(self, latents: torch.Tensor) -> torch.Tensor: ... + + def unpatchify(self, latents: torch.Tensor, output_shape: AudioLatentShape | VideoLatentShape) -> torch.Tensor: ... + + def get_token_count(self, target_shape: AudioLatentShape | VideoLatentShape) -> int: ... + + def get_patch_grid_bounds( + self, + output_shape: AudioLatentShape | VideoLatentShape, + device: torch.device | None = None, + ) -> torch.Tensor: ... + + +class VideoLatentPatchifier: + def __init__(self, patch_size: int): + self.patch_size = (1, patch_size, patch_size) + + def get_token_count(self, target_shape: VideoLatentShape) -> int: + return target_shape.frames * target_shape.height * target_shape.width // math.prod(self.patch_size) + + def patchify(self, latents: torch.Tensor) -> torch.Tensor: + return einops.rearrange( + latents, + "b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)", + p1=self.patch_size[0], + p2=self.patch_size[1], + p3=self.patch_size[2], + ) + + def unpatchify(self, latents: torch.Tensor, output_shape: VideoLatentShape) -> torch.Tensor: + return einops.rearrange( + latents, + "b (f h w) (c p q) -> b c f (h p) (w q)", + f=output_shape.frames // self.patch_size[0], + h=output_shape.height // self.patch_size[1], + w=output_shape.width // self.patch_size[2], + p=self.patch_size[1], + q=self.patch_size[2], + ) + + def get_patch_grid_bounds( + self, + output_shape: AudioLatentShape | VideoLatentShape, + device: torch.device | None = None, + ) -> torch.Tensor: + if not isinstance(output_shape, VideoLatentShape): + raise ValueError("VideoLatentPatchifier expects VideoLatentShape when computing coordinates") + grid_coords = torch.meshgrid( + torch.arange(start=0, end=output_shape.frames, step=self.patch_size[0], device=device), + torch.arange(start=0, end=output_shape.height, step=self.patch_size[1], device=device), + torch.arange(start=0, end=output_shape.width, step=self.patch_size[2], device=device), + indexing="ij", + ) + patch_starts = torch.stack(grid_coords, dim=0) + patch_size_delta = torch.tensor( + self.patch_size, + device=patch_starts.device, + dtype=patch_starts.dtype, + ).view(3, 1, 1, 1) + patch_ends = patch_starts + patch_size_delta + latent_coords = torch.stack((patch_starts, patch_ends), dim=-1) + return einops.repeat(latent_coords, "c f h w bounds -> b c (f h w) bounds", b=output_shape.batch, bounds=2) + + +class AudioPatchifier: + def __init__(self, patch_size: int): + # Keep the same latent layout as upstream LTX: + # audio latents are shaped (B, C, T, F) and patchify flattens along time -> (B, T, C*F). + # Positions encode real time in seconds so RoPE max_pos can be expressed in seconds. + self.patch_size = (patch_size, 1, 1) + self.sample_rate = 16000 + self.hop_length = 160 + self.audio_latent_downsample_factor = 4 + self.is_causal = True + self.shift = 0 + + def get_token_count(self, target_shape: AudioLatentShape) -> int: + return target_shape.frames // self.patch_size[0] + + def patchify(self, latents: torch.Tensor) -> torch.Tensor: + return einops.rearrange(latents, "b c (f p) m -> b f (c p m)", p=self.patch_size[0]) + + def unpatchify(self, latents: torch.Tensor, output_shape: AudioLatentShape) -> torch.Tensor: + return einops.rearrange( + latents, + "b f (c p m) -> b c (f p) m", + c=output_shape.channels, + p=self.patch_size[0], + m=output_shape.mel_bins, + ) + + def get_patch_grid_bounds( + self, + output_shape: AudioLatentShape | VideoLatentShape, + device: torch.device | None = None, + ) -> torch.Tensor: + if isinstance(output_shape, VideoLatentShape): + raise ValueError("AudioPatchifier expects AudioLatentShape when computing coordinates") + if device is None: + device = torch.device("cpu") + + start_latent = self.shift + end_latent = output_shape.frames + self.shift + audio_latent_frame_start = torch.arange(start_latent, end_latent, dtype=torch.float32, device=device) + audio_latent_frame_end = torch.arange(start_latent + 1, end_latent + 1, dtype=torch.float32, device=device) + + downsample = float(self.audio_latent_downsample_factor) + audio_mel_frame_start = audio_latent_frame_start * downsample + audio_mel_frame_end = audio_latent_frame_end * downsample + + if self.is_causal: + causal_offset = 1.0 + audio_mel_frame_start = (audio_mel_frame_start + causal_offset - downsample).clamp_min(0.0) + audio_mel_frame_end = (audio_mel_frame_end + causal_offset - downsample).clamp_min(0.0) + + start_timings = audio_mel_frame_start * float(self.hop_length) / float(self.sample_rate) + end_timings = audio_mel_frame_end * float(self.hop_length) / float(self.sample_rate) + + start_timings = start_timings.unsqueeze(0).expand(output_shape.batch, -1).unsqueeze(1) + end_timings = end_timings.unsqueeze(0).expand(output_shape.batch, -1).unsqueeze(1) + return torch.stack([start_timings, end_timings], dim=-1) + + +def get_pixel_coords( + latent_coords: torch.Tensor, + scale_factors: SpatioTemporalScaleFactors, + causal_fix: bool = False, +) -> torch.Tensor: + pixel_coords = latent_coords.clone() + pixel_coords[:, 0] *= scale_factors.time + pixel_coords[:, 1] *= scale_factors.height + pixel_coords[:, 2] *= scale_factors.width + if causal_fix: + pixel_coords[:, 0, :, :] -= scale_factors.time - 1 + pixel_coords[:, 0, :, :] = pixel_coords[:, 0, :, :].clamp_min(0) + return pixel_coords + + +@dataclass(frozen=True) +class LatentTools: + patchifier: Patchifier + target_shape: AudioLatentShape | VideoLatentShape + + def patchify(self, latent_state: LatentState) -> LatentState: + latent_state = latent_state.clone() + return replace( + latent_state, + latent=self.patchifier.patchify(latent_state.latent), + denoise_mask=self.patchifier.patchify(latent_state.denoise_mask), + clean_latent=self.patchifier.patchify(latent_state.clean_latent), + keyframes_mask=( + self.patchifier.patchify(latent_state.keyframes_mask) + if latent_state.keyframes_mask is not None + else None + ), + ) + + def unpatchify(self, latent_state: LatentState) -> LatentState: + latent_state = latent_state.clone() + return replace( + latent_state, + latent=self.patchifier.unpatchify(latent_state.latent, output_shape=self.target_shape), + denoise_mask=self.patchifier.unpatchify( + latent_state.denoise_mask, + output_shape=self.target_shape.mask_shape(), + ), + clean_latent=self.patchifier.unpatchify(latent_state.clean_latent, output_shape=self.target_shape), + keyframes_mask=( + self.patchifier.unpatchify(latent_state.keyframes_mask, output_shape=self.target_shape.mask_shape()) + if latent_state.keyframes_mask is not None + else None + ), + ) + + def clear_conditioning(self, latent_state: LatentState) -> LatentState: + num_tokens = self.patchifier.get_token_count(self.target_shape) + return LatentState( + latent=latent_state.latent[:, :num_tokens], + denoise_mask=torch.ones_like(latent_state.denoise_mask)[:, :num_tokens], + positions=latent_state.positions[:, :, :num_tokens], + clean_latent=latent_state.clean_latent[:, :num_tokens], + attention_mask=None, + keyframes_mask=( + latent_state.keyframes_mask[:, :num_tokens] if latent_state.keyframes_mask is not None else None + ), + ) + + +@dataclass(frozen=True) +class VideoLatentTools(LatentTools): + patchifier: VideoLatentPatchifier + target_shape: VideoLatentShape + fps: float + scale_factors: SpatioTemporalScaleFactors = VIDEO_SCALE_FACTORS + causal_fix: bool = True + + def create_initial_state( + self, + device: DeviceLikeType, + dtype: torch.dtype, + initial_latent: torch.Tensor | None = None, + ) -> LatentState: + if initial_latent is None: + initial_latent = torch.zeros(*self.target_shape.to_torch_shape(), device=device, dtype=dtype) + else: + initial_latent = initial_latent.to(device=device, dtype=dtype) + denoise_mask = torch.ones(*self.target_shape.mask_shape().to_torch_shape(), device=device, dtype=torch.float32) + latent_coords = self.patchifier.get_patch_grid_bounds(output_shape=self.target_shape, device=device) + positions = get_pixel_coords( + latent_coords, + self.scale_factors, + causal_fix=self.causal_fix, + ).to(dtype=torch.float32) + positions[:, 0, ...] /= self.fps + state = self.patchify( + LatentState( + latent=initial_latent, + denoise_mask=denoise_mask, + positions=positions, + clean_latent=initial_latent.clone(), + keyframes_mask=torch.zeros_like(denoise_mask), + ) + ) + assert state.keyframes_mask is not None + first_frame_tokens = self.patchifier.get_token_count(self.target_shape._replace(frames=1)) + keyframes_mask = state.keyframes_mask.clone() + keyframes_mask[:, :first_frame_tokens] = 1.0 + return replace(state, keyframes_mask=keyframes_mask) + + +@dataclass(frozen=True) +class AudioLatentTools(LatentTools): + patchifier: AudioPatchifier + target_shape: AudioLatentShape + + def create_initial_state( + self, + device: DeviceLikeType, + dtype: torch.dtype, + initial_latent: torch.Tensor | None = None, + ) -> LatentState: + if initial_latent is None: + initial_latent = torch.zeros(*self.target_shape.to_torch_shape(), device=device, dtype=dtype) + else: + initial_latent = initial_latent.to(device=device, dtype=dtype) + denoise_mask = torch.ones(*self.target_shape.mask_shape().to_torch_shape(), device=device, dtype=torch.float32) + return self.patchify( + LatentState( + latent=initial_latent, + denoise_mask=denoise_mask, + positions=self.patchifier.get_patch_grid_bounds( + output_shape=self.target_shape, + device=device, + ).to(dtype=torch.float32), + clean_latent=initial_latent.clone(), + ) + ) + + +class ConditioningError(RuntimeError): + pass + + +class ConditioningItem(Protocol): + def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState: ... + + +class VideoConditionByLatentIndex: + def __init__(self, latent: torch.Tensor, strength: float, latent_idx: int): + self.latent = latent + self.strength = strength + self.latent_idx = latent_idx + + def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState: + cond_batch, cond_channels, _, cond_height, cond_width = self.latent.shape + target_shape = latent_tools.target_shape + tgt_batch, tgt_channels, tgt_frames, tgt_height, tgt_width = target_shape.to_torch_shape() + if (cond_batch, cond_channels, cond_height, cond_width) != (tgt_batch, tgt_channels, tgt_height, tgt_width): + raise ConditioningError( + f"Can't apply image conditioning item to latent with shape {target_shape}, expected shape is " + f"({tgt_batch}, {tgt_channels}, {tgt_frames}, {tgt_height}, {tgt_width})." + ) + tokens = latent_tools.patchifier.patchify(self.latent) + start_token = latent_tools.patchifier.get_token_count(target_shape._replace(frames=self.latent_idx)) + stop_token = start_token + tokens.shape[1] + latent_state = latent_state.clone() + latent_state.latent[:, start_token:stop_token] = tokens + latent_state.clean_latent[:, start_token:stop_token] = tokens + latent_state.denoise_mask[:, start_token:stop_token] = 1.0 - self.strength + return latent_state + + +class VideoConditionByKeyframeIndex: + def __init__(self, keyframes: torch.Tensor, frame_idx: int, strength: float): + self.keyframes = keyframes + self.frame_idx = frame_idx + self.strength = strength + + def apply_to(self, latent_state: LatentState, latent_tools: VideoLatentTools) -> LatentState: + tokens = latent_tools.patchifier.patchify(self.keyframes) + positions = get_pixel_coords( + latent_coords=latent_tools.patchifier.get_patch_grid_bounds( + output_shape=VideoLatentShape.from_torch_shape(self.keyframes.shape), + device=self.keyframes.device, + ), + scale_factors=latent_tools.scale_factors, + causal_fix=latent_tools.causal_fix if self.frame_idx == 0 else False, + ).to(dtype=torch.float32) + positions[:, 0, ...] += self.frame_idx + positions[:, 0, ...] /= latent_tools.fps + denoise_mask = torch.full( + size=(*tokens.shape[:2], 1), + fill_value=1.0 - self.strength, + device=self.keyframes.device, + dtype=self.keyframes.dtype, + ) + return LatentState( + latent=torch.cat([latent_state.latent, tokens], dim=1), + denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1), + positions=torch.cat([latent_state.positions, positions], dim=2), + clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1), + attention_mask=update_attention_mask( + latent_state=latent_state, + attention_mask=None, + num_noisy_tokens=latent_tools.target_shape.token_count(), + num_new_tokens=tokens.shape[1], + batch_size=tokens.shape[0], + device=self.keyframes.device, + dtype=self.keyframes.dtype, + ), + ) + + +class ConditioningItemAttentionStrengthWrapper: + def __init__(self, conditioning: ConditioningItem, attention_mask: float | torch.Tensor): + self.conditioning = conditioning + self.attention_mask = attention_mask + + def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState: + original_state = latent_state + new_state = self.conditioning.apply_to(latent_state, latent_tools) + num_new_tokens = new_state.latent.shape[1] - original_state.latent.shape[1] + if num_new_tokens == 0: + return new_state + return replace( + new_state, + attention_mask=update_attention_mask( + latent_state=original_state, + attention_mask=self.attention_mask, + num_noisy_tokens=latent_tools.target_shape.token_count(), + num_new_tokens=num_new_tokens, + batch_size=new_state.latent.shape[0], + device=new_state.latent.device, + dtype=new_state.latent.dtype, + ), + ) + + +def resolve_cross_mask( + attention_mask: float | torch.Tensor, + num_new_tokens: int, + batch_size: int, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + if isinstance(attention_mask, float): + return torch.full((batch_size, num_new_tokens), attention_mask, device=device, dtype=dtype) + if attention_mask.ndim == 1: + return attention_mask[None].expand(batch_size, -1).to(device=device, dtype=dtype) + if attention_mask.ndim == 2: + return attention_mask.to(device=device, dtype=dtype) + raise ValueError(f"Unsupported attention mask shape: {attention_mask.shape}") + + +def build_attention_mask( + existing_mask: torch.Tensor | None, + num_noisy_tokens: int, + num_new_tokens: int, + num_existing_tokens: int, + cross_mask: torch.Tensor, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + batch_size = cross_mask.shape[0] + total_tokens = num_existing_tokens + num_new_tokens + attention_mask = torch.zeros((batch_size, total_tokens, total_tokens), device=device, dtype=dtype) + if existing_mask is not None: + attention_mask[:, :num_existing_tokens, :num_existing_tokens] = existing_mask + else: + attention_mask[:, :num_existing_tokens, :num_existing_tokens] = 1.0 + attention_mask[:, num_existing_tokens:, num_existing_tokens:] = 1.0 + attention_mask[:, :num_noisy_tokens, num_existing_tokens:] = cross_mask.unsqueeze(1) + attention_mask[:, num_existing_tokens:, :num_noisy_tokens] = cross_mask.unsqueeze(2) + return attention_mask + + +def update_attention_mask( + latent_state: LatentState, + attention_mask: float | torch.Tensor | None, + num_noisy_tokens: int, + num_new_tokens: int, + batch_size: int, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor | None: + if attention_mask is None: + if latent_state.attention_mask is None: + return None + cross_mask = torch.ones(batch_size, num_new_tokens, device=device, dtype=dtype) + return build_attention_mask( + existing_mask=latent_state.attention_mask, + num_noisy_tokens=num_noisy_tokens, + num_new_tokens=num_new_tokens, + num_existing_tokens=latent_state.latent.shape[1], + cross_mask=cross_mask, + device=device, + dtype=dtype, + ) + return build_attention_mask( + existing_mask=latent_state.attention_mask, + num_noisy_tokens=num_noisy_tokens, + num_new_tokens=num_new_tokens, + num_existing_tokens=latent_state.latent.shape[1], + cross_mask=resolve_cross_mask(attention_mask, num_new_tokens, batch_size, device, dtype), + device=device, + dtype=dtype, + ) + + +@dataclass(frozen=True) +class MultiModalGuiderParams: + cfg_scale: float = 1.0 + stg_scale: float = 0.0 + stg_blocks: list[int] | None = field(default_factory=list) + rescale_scale: float = 0.0 + modality_scale: float = 1.0 + skip_step: int = 0 + + +@dataclass(frozen=True) +class MultiModalGuider: + params: MultiModalGuiderParams + negative_context: torch.Tensor | None = None + + def calculate( + self, + cond: torch.Tensor, + uncond_text: torch.Tensor | float, + uncond_perturbed: torch.Tensor | float, + uncond_modality: torch.Tensor | float, + ) -> torch.Tensor: + pred = ( + cond + + (self.params.cfg_scale - 1) * (cond - uncond_text) + + self.params.stg_scale * (cond - uncond_perturbed) + + (self.params.modality_scale - 1) * (cond - uncond_modality) + ) + if self.params.rescale_scale != 0: + factor = cond.std() / pred.std() + factor = self.params.rescale_scale * factor + (1 - self.params.rescale_scale) + pred = pred * factor + return pred + + def do_unconditional_generation(self) -> bool: + return not math.isclose(self.params.cfg_scale, 1.0) + + def do_perturbed_generation(self) -> bool: + return not math.isclose(self.params.stg_scale, 0.0) + + def do_isolated_modality_generation(self) -> bool: + return not math.isclose(self.params.modality_scale, 1.0) + + def should_skip_step(self, step_index: int) -> bool: + if self.params.skip_step == 0: + return False + return step_index % (self.params.skip_step + 1) != 0 diff --git a/telefuser/pipelines/ltx25_distilled/latent_upsampling.py b/telefuser/pipelines/ltx25_distilled/latent_upsampling.py new file mode 100644 index 0000000..ecd5f76 --- /dev/null +++ b/telefuser/pipelines/ltx25_distilled/latent_upsampling.py @@ -0,0 +1,27 @@ +"""Latent resolution bridge for LTX-2.5.""" + +from __future__ import annotations + +import torch + +from telefuser.core.base_stage import BaseStage, with_model_offload +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.models.ltx25 import LTX25SpatialUpsampler +from telefuser.models.ltx25.spatial_upsampler import LTX25PerChannelStatistics + + +class LTX25LatentUpsamplingStage(BaseStage): + def __init__(self, module_manager: ModuleManager, config: ModelRuntimeConfig) -> None: + super().__init__("ltx25_latent_upsampling", config) + self.spatial_upsampler: LTX25SpatialUpsampler = module_manager.fetch_module("ltx25_spatial_upsampler") + self.latent_statistics: LTX25PerChannelStatistics = module_manager.fetch_module("ltx25_video_latent_statistics") + self.model_names = ["spatial_upsampler", "latent_statistics"] + + @with_model_offload(["spatial_upsampler", "latent_statistics"]) + @torch.inference_mode() + def process(self, latent: torch.Tensor) -> torch.Tensor: + return self.latent_statistics.normalize(self.spatial_upsampler(self.latent_statistics.un_normalize(latent))) + + +__all__ = ["LTX25LatentUpsamplingStage"] diff --git a/telefuser/pipelines/ltx25_distilled/loader.py b/telefuser/pipelines/ltx25_distilled/loader.py new file mode 100644 index 0000000..893f7ea --- /dev/null +++ b/telefuser/pipelines/ltx25_distilled/loader.py @@ -0,0 +1,93 @@ +"""ModuleManager loading helpers for the LTX-2.5 distilled model pack.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Literal + +import torch + +from telefuser.core.module_manager import ModuleManager +from telefuser.models.ltx25 import ( + DiffusionVideoDecoder, + LTX25AVTransformer, + LTX25ConvVideoVAE, + LTX25DurationHead, + LTX25EmbeddingsProcessor, + LTX25Gemma4TextEncoder, + LTX25ModelPaths, + LTX25SpatialUpsampler, + LTX25VideoEncoder, + load_ltx25_audio_decoder_and_vocoder, +) +from telefuser.models.ltx25.spatial_upsampler import load_video_latent_statistics + + +def load_ltx25_distilled_modules( + module_manager: ModuleManager, + model_root: str | Path, + *, + video_vae: Literal["diff", "conv"] = "diff", + torch_dtype: torch.dtype = torch.bfloat16, +) -> LTX25ModelPaths: + """Load every LTX-2.5 component on CPU and register it with ``module_manager``.""" + if video_vae not in ("diff", "conv"): + raise ValueError(f"video_vae must be 'diff' or 'conv', got {video_vae!r}") + paths = LTX25ModelPaths.from_model_root(model_root) + video_vae_path = paths.video_vae_path if video_vae == "diff" else paths.conv_video_vae_path + + def add(module: torch.nn.Module, name: str, path: str | Path) -> None: + module_manager.add_module(module, name, path=str(path)) + + add( + LTX25Gemma4TextEncoder.from_checkpoint(paths.text_encoder_path, device="cpu", torch_dtype=torch_dtype), + "ltx25_gemma4", + paths.text_encoder_path, + ) + add( + LTX25EmbeddingsProcessor.from_checkpoints( + paths.transformer_path, + paths.text_encoder_path, + device="cpu", + torch_dtype=torch_dtype, + ), + "ltx25_embeddings_processor", + paths.transformer_path, + ) + add( + LTX25DurationHead.from_checkpoint(paths.duration_head_path, device="cpu", torch_dtype=torch_dtype), + "ltx25_duration_head", + paths.duration_head_path, + ) + add( + LTX25VideoEncoder.from_checkpoint(video_vae_path, device="cpu", torch_dtype=torch_dtype), + "ltx25_video_encoder", + video_vae_path, + ) + add( + LTX25AVTransformer.from_checkpoint(paths.transformer_path, device="cpu", torch_dtype=torch_dtype), + "ltx25_transformer", + paths.transformer_path, + ) + add( + LTX25SpatialUpsampler.from_checkpoint(paths.spatial_upsampler_path, device="cpu", torch_dtype=torch_dtype), + "ltx25_spatial_upsampler", + paths.spatial_upsampler_path, + ) + add(load_video_latent_statistics(video_vae_path), "ltx25_video_latent_statistics", video_vae_path) + + if video_vae == "diff": + video_decoder = DiffusionVideoDecoder.from_checkpoint(video_vae_path, device="cpu", torch_dtype=torch_dtype) + else: + video_decoder = LTX25ConvVideoVAE.from_checkpoint(video_vae_path, device="cpu", torch_dtype=torch_dtype) + add(video_decoder, "ltx25_video_decoder", video_vae_path) + + audio_decoder, vocoder = load_ltx25_audio_decoder_and_vocoder( + paths.audio_vae_path, device="cpu", torch_dtype=torch_dtype + ) + add(audio_decoder, "ltx25_audio_decoder", paths.audio_vae_path) + add(vocoder, "ltx25_vocoder", paths.audio_vae_path) + return paths + + +__all__ = ["load_ltx25_distilled_modules"] diff --git a/telefuser/pipelines/ltx25_distilled/pipeline.py b/telefuser/pipelines/ltx25_distilled/pipeline.py new file mode 100644 index 0000000..8598705 --- /dev/null +++ b/telefuser/pipelines/ltx25_distilled/pipeline.py @@ -0,0 +1,360 @@ +"""Multi-stage LTX-2.5 distilled pipeline runtime.""" + +from __future__ import annotations + +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Literal, Sequence + +import torch +from PIL import Image + +from telefuser.core.base_pipeline import BasePipeline +from telefuser.core.config import ( + AttentionConfig, + AttnImplType, + ModelRuntimeConfig, + OffloadConfig, + ParallelConfig, + WeightOffloadType, +) +from telefuser.core.module_manager import ModuleManager +from telefuser.models.ltx25.diff_vae.types import VideoLatentShape +from telefuser.models.ltx25.sampler import ANCESTRAL_NOISE_SEED_OFFSET, LTX25_STAGE2_DISTILLED_SIGMAS +from telefuser.utils.func import auto_async_call +from telefuser.worker.parallel_worker import ParallelWorker + +from .audio_decoding import LTX25AudioDecodingStage +from .denoising import LTX25DenoisingStage +from .latent import ( + AudioLatentShape, + AudioLatentTools, + AudioPatchifier, + LatentState, + VideoLatentPatchifier, + VideoLatentTools, +) +from .latent_upsampling import LTX25LatentUpsamplingStage +from .loader import load_ltx25_distilled_modules +from .text_encoding import LTX25TextEncodingStage +from .video_conditioning import LTX25VideoConditioningStage +from .video_decoding import LTX25VideoDecodingStage + + +@dataclass(frozen=True, slots=True) +class LTX25ImageCondition: + image: Image.Image + frame_idx: int = 0 + strength: float = 1.0 + crf: int | None = None + + +@dataclass(frozen=True, slots=True) +class LTX25DistilledOutput: + video_chunks: tuple[torch.Tensor, ...] + audio: torch.Tensor + video_latent: torch.Tensor + audio_latent: torch.Tensor + num_frames: int + frame_rate: float + + +@dataclass +class LTX25DistilledConfig: + """Runtime settings for independently managed LTX-2.5 stages.""" + + video_vae: Literal["diff", "conv"] = "diff" + text_encoding_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) + video_conditioning_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) + denoising_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) + upsampling_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) + video_decoding_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) + audio_decoding_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) + + +class LTX25DistilledPipeline(BasePipeline): + """Two-stage distilled pipeline composed from ModuleManager-backed stages.""" + + def __init__(self, device: str | torch.device = "cuda", torch_dtype: torch.dtype = torch.bfloat16) -> None: + super().__init__(device=device, torch_dtype=torch_dtype) + self.config: LTX25DistilledConfig | None = None + self.text_stage: LTX25TextEncodingStage | None = None + self.conditioning_stage: LTX25VideoConditioningStage | None = None + self.denoising_stage: LTX25DenoisingStage | ParallelWorker | None = None + self.upsampling_stage: LTX25LatentUpsamplingStage | None = None + self.video_decoding_stage: LTX25VideoDecodingStage | None = None + self.audio_decoding_stage: LTX25AudioDecodingStage | None = None + + def init(self, module_manager: ModuleManager, config: LTX25DistilledConfig) -> None: + self.config = config + self.text_stage = LTX25TextEncodingStage(module_manager, config.text_encoding_config) + self.conditioning_stage = LTX25VideoConditioningStage(module_manager, config.video_conditioning_config) + denoising_stage = LTX25DenoisingStage(module_manager, config.denoising_config) + self.denoising_stage = ( + ParallelWorker(denoising_stage) + if config.denoising_config.parallel_config.world_size > 1 + else denoising_stage + ) + self.upsampling_stage = LTX25LatentUpsamplingStage(module_manager, config.upsampling_config) + self.video_decoding_stage = LTX25VideoDecodingStage( + module_manager, config.video_decoding_config, video_vae=config.video_vae + ) + self.audio_decoding_stage = LTX25AudioDecodingStage(module_manager, config.audio_decoding_config) + self._model_info = module_manager.get_model_info() + + @classmethod + def from_model_root( + cls, + model_root: str | Path, + *, + device: str = "cuda", + torch_dtype: torch.dtype = torch.bfloat16, + video_vae: Literal["diff", "conv"] = "diff", + offload: Literal["none", "cpu"] = "cpu", + parallelism: int = 1, + attn_impl: AttnImplType = AttnImplType.FLASH_ATTN_4, + ) -> "LTX25DistilledPipeline": + module_manager = ModuleManager(device="cpu", torch_dtype=torch_dtype) + load_ltx25_distilled_modules(module_manager, model_root, video_vae=video_vae, torch_dtype=torch_dtype) + pipeline = cls(device=device, torch_dtype=torch_dtype) + pipeline.init( + module_manager, + build_ltx25_distilled_config( + device, + torch_dtype, + video_vae, + offload, + parallelism=parallelism, + attn_impl=attn_impl, + ), + ) + return pipeline + + def close(self) -> None: + """Release distributed denoising workers, when configured.""" + if isinstance(self.denoising_stage, ParallelWorker): + self.denoising_stage.close() + + def __del__(self) -> None: + try: + self.close() + except Exception: + pass + + def _denoising_input(self, state: LatentState) -> LatentState | dict[str, torch.Tensor | None]: + if not isinstance(self.denoising_stage, ParallelWorker): + return state + return { + "latent": state.latent, + "denoise_mask": state.denoise_mask, + "positions": state.positions, + "clean_latent": state.clean_latent, + "attention_mask": state.attention_mask, + "keyframes_mask": state.keyframes_mask, + } + + def _get_stages(self) -> list[object]: + return [ + stage + for stage in ( + self.text_stage, + self.conditioning_stage, + self.denoising_stage, + self.upsampling_stage, + self.video_decoding_stage, + self.audio_decoding_stage, + ) + if stage is not None + ] + + @torch.inference_mode() + def __call__( + self, + prompt: str, + *, + seed: int, + height: int, + width: int, + num_frames: int | None = None, + frame_rate: float = 24.0, + images: Sequence[LTX25ImageCondition] = (), + ) -> LTX25DistilledOutput: + if ( + self.config is None + or self.text_stage is None + or self.conditioning_stage is None + or self.denoising_stage is None + or self.upsampling_stage is None + or self.video_decoding_stage is None + or self.audio_decoding_stage is None + ): + raise RuntimeError("LTX25DistilledPipeline.init must be called before generation") + + _validate_resolution(height, width, frame_rate) + generator = torch.Generator(device=self.device).manual_seed(seed) + video_context, audio_context = self.text_stage.encode(prompt) + if num_frames is None: + num_frames = self.text_stage.predict_num_frames(video_context, audio_context, frame_rate) + _validate_request(height, width, num_frames, frame_rate) + + stage_1_tools = _video_tools( + batch=1, frames=num_frames, height=height // 2, width=width // 2, frame_rate=frame_rate + ) + audio_tools = _audio_tools(num_frames=num_frames, frame_rate=frame_rate) + stage_1_video = stage_1_tools.create_initial_state(self.device, self.torch_dtype) + if images: + stage_1_video = self.conditioning_stage.apply(stage_1_video, stage_1_tools, images, height // 2, width // 2) + stage_1_video = _noised_state(stage_1_video, 1.0, generator) + stage_1_audio = _noised_state(audio_tools.create_initial_state(self.device, self.torch_dtype), 1.0, generator) + + try: + stage_1_video, stage_1_audio = auto_async_call( + self.denoising_stage.denoise_stage1, + self._denoising_input(stage_1_video), + self._denoising_input(stage_1_audio), + video_context, + audio_context, + noise_seed=seed + ANCESTRAL_NOISE_SEED_OFFSET, + )() + low_resolution_latent = stage_1_tools.unpatchify(stage_1_tools.clear_conditioning(stage_1_video)).latent + upscaled_video = self.upsampling_stage.process(low_resolution_latent) + + stage_2_tools = _video_tools(batch=1, frames=num_frames, height=height, width=width, frame_rate=frame_rate) + stage_2_video = stage_2_tools.create_initial_state(self.device, self.torch_dtype, upscaled_video) + if images: + stage_2_video = self.conditioning_stage.apply(stage_2_video, stage_2_tools, images, height, width) + stage_2_video = _noised_state(stage_2_video, LTX25_STAGE2_DISTILLED_SIGMAS[0], generator) + stage_2_audio = _noised_state( + audio_tools.create_initial_state( + self.device, self.torch_dtype, audio_tools.unpatchify(stage_1_audio).latent + ), + LTX25_STAGE2_DISTILLED_SIGMAS[0], + generator, + ) + stage_2_video, stage_2_audio = auto_async_call( + self.denoising_stage.denoise_stage2, + self._denoising_input(stage_2_video), + self._denoising_input(stage_2_audio), + video_context, + audio_context, + )() + finally: + if not isinstance(self.denoising_stage, ParallelWorker) or not self.denoising_stage.failed: + auto_async_call(self.denoising_stage.finish_request)() + + video_latent = stage_2_tools.unpatchify(stage_2_tools.clear_conditioning(stage_2_video)).latent + audio_latent = audio_tools.unpatchify(stage_2_audio).latent + return LTX25DistilledOutput( + self.video_decoding_stage.decode(video_latent, generator), + self.audio_decoding_stage.decode(audio_latent), + video_latent, + audio_latent, + num_frames, + frame_rate, + ) + + +def build_ltx25_distilled_config( + device: str, + torch_dtype: torch.dtype, + video_vae: Literal["diff", "conv"], + offload: Literal["none", "cpu"], + *, + parallelism: int = 1, + attn_impl: AttnImplType = AttnImplType.FLASH_ATTN_4, +) -> LTX25DistilledConfig: + if video_vae not in ("diff", "conv"): + raise ValueError(f"video_vae must be 'diff' or 'conv', got {video_vae!r}") + if offload not in ("none", "cpu"): + raise ValueError(f"offload must be 'none' or 'cpu', got {offload!r}") + if parallelism < 1 or 32 % parallelism: + raise ValueError(f"parallelism must be a positive divisor of 32, got {parallelism}") + if not isinstance(attn_impl, AttnImplType): + raise TypeError(f"attn_impl must be an AttnImplType, got {type(attn_impl).__name__}") + attention_config = AttentionConfig.dense_attention(attn_impl) + if attention_config.is_sparse(): + raise ValueError(f"LTX-2.5 supports dense attention implementations only, got {attn_impl.name}") + target_device = torch.device(device) + device_type = target_device.type + device_id = 0 if target_device.index is None else target_device.index + regular = WeightOffloadType.MODEL_CPU_OFFLOAD if offload == "cpu" else WeightOffloadType.NO_CPU_OFFLOAD + denoising = ( + WeightOffloadType.ASYNC_CPU_OFFLOAD + if offload == "cpu" and parallelism == 1 + else WeightOffloadType.NO_CPU_OFFLOAD + ) + + def runtime(offload_type: WeightOffloadType, *, distributed: bool = False) -> ModelRuntimeConfig: + return ModelRuntimeConfig( + device_type=device_type, + device_id=device_id, + torch_dtype=torch_dtype, + offload_config=OffloadConfig(offload_type=offload_type), + attention_config=attention_config, + parallel_config=( + ParallelConfig( + device_ids=list(range(parallelism)), + sp_ulysses_degree=parallelism, + enable_fsdp=True, + timeout=1800, + ) + if distributed + else ParallelConfig() + ), + ) + + return LTX25DistilledConfig( + video_vae=video_vae, + text_encoding_config=runtime(regular), + video_conditioning_config=runtime(regular), + denoising_config=runtime(denoising, distributed=parallelism > 1), + upsampling_config=runtime(regular), + video_decoding_config=runtime(regular), + audio_decoding_config=runtime(regular), + ) + + +def _video_tools(*, batch: int, frames: int, height: int, width: int, frame_rate: float) -> VideoLatentTools: + shape = VideoLatentShape( + batch=batch, + channels=128, + frames=(frames - 1) // 8 + 1, + height=height // 32, + width=width // 32, + ) + return VideoLatentTools(VideoLatentPatchifier(1), shape, frame_rate) + + +def _audio_tools(*, num_frames: int, frame_rate: float) -> AudioLatentTools: + return AudioLatentTools( + AudioPatchifier(1), AudioLatentShape.from_duration(batch=1, duration=num_frames / frame_rate) + ) + + +def _noised_state(state: LatentState, noise_scale: float, generator: torch.Generator) -> LatentState: + noise = torch.randn(state.latent.shape, generator=generator, dtype=state.latent.dtype, device=state.latent.device) + latent = torch.lerp(state.latent.float(), noise.float(), noise_scale) + latent = torch.lerp(state.clean_latent.float(), latent, state.denoise_mask) + return replace(state, latent=latent.to(state.latent.dtype)) + + +def _validate_resolution(height: int, width: int, frame_rate: float) -> None: + if height <= 0 or width <= 0 or height % 64 or width % 64: + raise ValueError("LTX-2.5 distilled height and width must be positive multiples of 64") + if frame_rate <= 0: + raise ValueError("frame_rate must be positive") + + +def _validate_request(height: int, width: int, num_frames: int, frame_rate: float) -> None: + _validate_resolution(height, width, frame_rate) + if num_frames < 1 or (num_frames - 1) % 8: + raise ValueError("LTX-2.5 num_frames must satisfy num_frames = 8k + 1") + + +__all__ = [ + "LTX25DistilledConfig", + "LTX25DistilledOutput", + "LTX25DistilledPipeline", + "LTX25ImageCondition", + "build_ltx25_distilled_config", +] diff --git a/telefuser/pipelines/ltx25_distilled/reference.py b/telefuser/pipelines/ltx25_distilled/reference.py new file mode 100644 index 0000000..a00b11f --- /dev/null +++ b/telefuser/pipelines/ltx25_distilled/reference.py @@ -0,0 +1,588 @@ +"""Faithful monolithic LTX-2.5 distilled text-to-video diffusion reference path.""" + +from __future__ import annotations + +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Callable, Literal, Protocol, Sequence + +import torch +from PIL import Image + +from telefuser.models.ltx25.checkpoint import LTX25ModelPaths, inspect_checkpoint +from telefuser.models.ltx25.diff_vae.types import VideoLatentShape, VideoPixelShape +from telefuser.models.ltx25.embeddings import LTX25EmbeddingsProcessorOutput +from telefuser.models.ltx25.sampler import LTX25EulerAncestralStep, ancestral_noise_generator, distilled_sigmas +from telefuser.models.ltx25.transformer import BatchedPerturbationConfig, Modality + +from .image import default_image_crf, preprocess_ltx25_image +from .latent import ( + AudioLatentShape, + AudioLatentTools, + AudioPatchifier, + LatentState, + VideoConditionByKeyframeIndex, + VideoConditionByLatentIndex, + VideoLatentPatchifier, + VideoLatentTools, +) + + +def _release_modules(*modules: object) -> None: + """Return lazily loaded checkpoint modules to CPU between reference phases.""" + for module in modules: + if isinstance(module, (_LazyTextEncoder, _LazyCallable)): + module.release() + elif isinstance(module, torch.nn.Module): + module.to("cpu") + if torch.cuda.is_available(): + torch.cuda.synchronize() + torch.cuda.empty_cache() + + +class TextEncoder(Protocol): + """Minimal Gemma interface consumed by the distilled reference path.""" + + def encode(self, prompts: list[str]) -> tuple[tuple[torch.Tensor, ...], torch.Tensor, torch.Tensor]: ... + + +class EmbeddingsProcessor(Protocol): + """Minimal dual-context connector interface consumed by the reference path.""" + + def __call__( + self, hidden_states: tuple[torch.Tensor, ...], attention_mask: torch.Tensor + ) -> LTX25EmbeddingsProcessorOutput: ... + + +class X0Transformer(Protocol): + """Joint video/audio x0 prediction interface.""" + + def __call__( + self, + video: Modality | None, + audio: Modality | None, + perturbations: BatchedPerturbationConfig, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: ... + + +class SpatialUpsampler(Protocol): + """Spatial latent upsampler interface.""" + + def __call__(self, latent: torch.Tensor) -> torch.Tensor: ... + + +class LatentStatistics(Protocol): + """Video-VAE latent normalization bridge used by the learned upsampler.""" + + def normalize(self, latent: torch.Tensor) -> torch.Tensor: ... + + def un_normalize(self, latent: torch.Tensor) -> torch.Tensor: ... + + +class VideoEncoder(Protocol): + """Video VAE encoder used to build an I2V latent condition.""" + + def __call__(self, pixels: torch.Tensor) -> torch.Tensor: ... + + +class _LazyTextEncoder: + def __init__(self, loader: Callable[[], TextEncoder], *, release_after_call: bool = True) -> None: + self._loader = loader + self._release_after_call = release_after_call + self._model: TextEncoder | None = None + + def encode(self, prompts: list[str]) -> tuple[tuple[torch.Tensor, ...], torch.Tensor, torch.Tensor]: + if self._model is None: + self._model = self._loader() + try: + return self._model.encode(prompts) + finally: + if self._release_after_call: + _release_modules(self._model) + self._model = None + + def release(self) -> None: + """Release a retained text encoder after the reference path no longer needs it.""" + _release_modules(self._model) + self._model = None + + +class _LazyCallable: + def __init__(self, loader: Callable[[], Callable[..., object]], *, release_after_call: bool = False) -> None: + self._loader = loader + self._release_after_call = release_after_call + self._model: object | None = None + + def resolve(self) -> object: + if self._model is None: + self._model = self._loader() + return self._model + + def release(self) -> None: + """Release a retained callable checkpoint model.""" + _release_modules(self._model) + self._model = None + + def __getattr__(self, name: str) -> object: + """Proxy model attributes required by instrumentation before the first call.""" + return getattr(self.resolve(), name) + + def forward(self, *args: object, **kwargs: object) -> object: + self.resolve() + try: + return self._model(*args, **kwargs) # type: ignore[operator] + finally: + if self._release_after_call: + _release_modules(self._model) + self._model = None + + def __call__(self, *args: object, **kwargs: object) -> object: + return self.forward(*args, **kwargs) + + +@dataclass(frozen=True, slots=True) +class LTX25ReferenceImageCondition: + """An I2V condition in output-frame coordinates for Golden capture.""" + + image: Image.Image + frame_idx: int = 0 + strength: float = 1.0 + crf: int | None = None + + +@dataclass(frozen=True, slots=True) +class LTX25ReferenceRequest: + """Fixed-shape T2V/I2V request for the faithful pre-stage-splitting path.""" + + prompt: str + seed: int + height: int + width: int + num_frames: int + frame_rate: float + images: tuple[LTX25ReferenceImageCondition, ...] = () + + def validate(self) -> None: + if self.height <= 0 or self.width <= 0 or self.num_frames <= 0 or self.frame_rate <= 0: + raise ValueError("height, width, num_frames, and frame_rate must be positive") + if self.height % 64 or self.width % 64: + raise ValueError("LTX-2.5 distilled two-stage generation requires height and width divisible by 64") + for condition in self.images: + if condition.frame_idx < 0: + raise ValueError(f"image frame_idx must be non-negative, got {condition.frame_idx}") + if not 0.0 <= condition.strength <= 1.0: + raise ValueError(f"image strength must be in [0, 1], got {condition.strength}") + + +@dataclass(frozen=True, slots=True) +class LTX25ReferenceComponents: + """Already-loaded modules required by the monolithic diffusion reference path.""" + + text_encoder: TextEncoder + embeddings_processor: EmbeddingsProcessor + transformer: X0Transformer + spatial_upsampler: SpatialUpsampler + latent_statistics: LatentStatistics + video_encoder: VideoEncoder | None = None + + +@dataclass(slots=True) +class LTX25ReferenceTrace: + """Intermediate tensors used for golden-artifact comparison.""" + + video_context: torch.Tensor + audio_context: torch.Tensor + context_attention_mask: torch.Tensor + artifacts: dict[str, torch.Tensor] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True) +class LTX25ReferenceResult: + """Final unpatchified stage states, decoder RNG state, and comparison trace.""" + + stage1_video: LatentState + stage1_audio: LatentState + stage2_video: LatentState + stage2_audio: LatentState + decoder_generator_state: torch.Tensor + trace: LTX25ReferenceTrace + + +def _noised_state(state: LatentState, generator: torch.Generator, noise_scale: float) -> LatentState: + """Match upstream GaussianNoiser's float32 lerp and original latent dtype.""" + noise = torch.randn( + *state.latent.shape, + device=state.latent.device, + dtype=state.latent.dtype, + generator=generator, + ) + latent = torch.lerp(state.latent.float(), noise.float(), noise_scale) + latent = torch.lerp(state.clean_latent.float(), latent, state.denoise_mask) + return replace(state, latent=latent.to(state.latent.dtype)) + + +def _modality(state: LatentState, context: torch.Tensor, sigma: torch.Tensor) -> Modality: + """Construct the upstream SimpleDenoiser modality without guidance or conditioning masks.""" + batch_size = state.latent.shape[0] + expanded_sigma = sigma.expand(batch_size) + return Modality( + latent=state.latent, + sigma=expanded_sigma, + timesteps=state.denoise_mask * expanded_sigma.view(batch_size, 1, 1), + positions=state.positions, + context=context, + context_mask=None, + attention_mask=state.attention_mask, + keyframes_mask=state.keyframes_mask, + ) + + +def _post_process(denoised: torch.Tensor, state: LatentState) -> torch.Tensor: + return (denoised * state.denoise_mask + state.clean_latent.float() * (1 - state.denoise_mask)).to(denoised.dtype) + + +def _deterministic_euler_step( + sample: torch.Tensor, + denoised: torch.Tensor, + sigmas: torch.Tensor, + step_index: int, +) -> torch.Tensor: + """Match the upstream EulerDiffusionStep rounding points.""" + sigma = sigmas[step_index] + velocity = ((sample.float() - denoised.float()) / sigma.to(torch.float32).item()).to(sample.dtype) + return (sample.float() + velocity.float() * (sigmas[step_index + 1] - sigma)).to(sample.dtype) + + +class LTX25DistilledReference: + """Faithful two-stage T2V diffusion path used to establish upstream parity.""" + + def __init__( + self, + components: LTX25ReferenceComponents, + *, + device: torch.device, + dtype: torch.dtype, + capture_prompt_intermediates: bool = False, + video_encoder_path: str | Path | None = None, + offload: Literal["none", "cpu"] = "cpu", + ) -> None: + self.components = components + self.device = device + self.dtype = dtype + self.capture_prompt_intermediates = capture_prompt_intermediates + self.video_encoder_path = Path(video_encoder_path) if video_encoder_path is not None else None + self.offload = offload + + @classmethod + def from_model_root( + cls, + model_root: str, + *, + device: torch.device | str = "cuda", + dtype: torch.dtype = torch.bfloat16, + video_vae: Literal["diff", "conv"] = "diff", + capture_prompt_intermediates: bool = False, + offload: Literal["none", "cpu"] = "cpu", + ) -> "LTX25DistilledReference": + """Load the exact modules needed by the monolithic T2V parity path.""" + from telefuser.models.ltx25.embeddings import LTX25EmbeddingsProcessor + from telefuser.models.ltx25.gemma4 import LTX25Gemma4TextEncoder + from telefuser.models.ltx25.spatial_upsampler import LTX25SpatialUpsampler, load_video_latent_statistics + from telefuser.models.ltx25.transformer import LTX25AVTransformer + + resolved_device = torch.device(device) + paths = LTX25ModelPaths.from_model_root(model_root) + video_vae_path = paths.video_vae_path if video_vae == "diff" else paths.conv_video_vae_path + components = LTX25ReferenceComponents( + text_encoder=_LazyTextEncoder( + lambda: LTX25Gemma4TextEncoder.from_checkpoint( + paths.text_encoder_path, device=resolved_device, torch_dtype=dtype + ), + release_after_call=offload == "cpu", + ), + embeddings_processor=_LazyCallable( + lambda: LTX25EmbeddingsProcessor.from_checkpoints( + paths.transformer_path, paths.text_encoder_path, device=resolved_device, torch_dtype=dtype + ), + release_after_call=offload == "cpu", + ), + transformer=_LazyCallable( + lambda: LTX25AVTransformer.from_checkpoint( + paths.transformer_path, device=resolved_device, torch_dtype=dtype + ) + ), + spatial_upsampler=_LazyCallable( + lambda: LTX25SpatialUpsampler.from_checkpoint( + paths.spatial_upsampler_path, device=resolved_device, torch_dtype=dtype + ), + release_after_call=offload == "cpu", + ), + latent_statistics=load_video_latent_statistics(video_vae_path).to(device=resolved_device, dtype=dtype), + ) + return cls( + components, + device=resolved_device, + dtype=dtype, + capture_prompt_intermediates=capture_prompt_intermediates, + video_encoder_path=video_vae_path, + offload=offload, + ) + + @torch.inference_mode() + def generate(self, request: LTX25ReferenceRequest) -> LTX25ReferenceResult: + """Run the upstream-equivalent T2V/I2V diffusion and return unpatchified stage states.""" + request.validate() + contexts, prompt_artifacts = self._encode_prompt(request.prompt) + trace = LTX25ReferenceTrace( + video_context=contexts.video_encoding.detach().clone(), + audio_context=contexts.audio_encoding.detach().clone(), + context_attention_mask=contexts.attention_mask.detach().clone(), + artifacts=prompt_artifacts, + ) + generator = torch.Generator(device=self.device).manual_seed(request.seed) + + stage1_video_tools, stage1_audio_tools = self._tools(request, half_resolution=True) + stage1_video = self._apply_image_conditions( + stage1_video_tools.create_initial_state(self.device, self.dtype), + stage1_video_tools, + request.images, + request.height // 2, + request.width // 2, + ) + stage1_video = _noised_state(stage1_video, generator, 1.0) + stage1_audio = _noised_state(stage1_audio_tools.create_initial_state(self.device, self.dtype), generator, 1.0) + stage1_video, stage1_audio = self._sample_stage( + stage_name="stage1", + video=stage1_video, + audio=stage1_audio, + video_context=contexts.video_encoding, + audio_context=contexts.audio_encoding, + sigmas=distilled_sigmas(1, device=self.device), + ancestral=True, + seed=request.seed, + trace=trace, + ) + stage1_video = stage1_video_tools.unpatchify(stage1_video) + stage1_audio = stage1_audio_tools.unpatchify(stage1_audio) + + upsampled = self.components.latent_statistics.normalize( + self.components.spatial_upsampler(self.components.latent_statistics.un_normalize(stage1_video.latent[:1])) + ) + trace.artifacts["upsampler_input"] = stage1_video.latent.detach().clone() + trace.artifacts["upsampler_output"] = upsampled.detach().clone() + stage2_video_tools, stage2_audio_tools = self._tools(request, half_resolution=False) + stage2_sigmas = distilled_sigmas(2, device=self.device) + stage2_video = self._apply_image_conditions( + stage2_video_tools.create_initial_state(self.device, self.dtype, initial_latent=upsampled), + stage2_video_tools, + request.images, + request.height, + request.width, + ) + stage2_video = _noised_state( + stage2_video, + generator, + float(stage2_sigmas[0]), + ) + stage2_audio = _noised_state( + stage2_audio_tools.create_initial_state(self.device, self.dtype, initial_latent=stage1_audio.latent), + generator, + float(stage2_sigmas[0]), + ) + stage2_video, stage2_audio = self._sample_stage( + stage_name="stage2", + video=stage2_video, + audio=stage2_audio, + video_context=contexts.video_encoding, + audio_context=contexts.audio_encoding, + sigmas=stage2_sigmas, + ancestral=False, + seed=request.seed, + trace=trace, + ) + stage2_video = stage2_video_tools.unpatchify(stage2_video) + stage2_audio = stage2_audio_tools.unpatchify(stage2_audio) + trace.artifacts["stage1_video_latent"] = stage1_video.latent.detach().clone() + trace.artifacts["stage1_audio_latent"] = stage1_audio.latent.detach().clone() + trace.artifacts["stage2_video_latent"] = stage2_video.latent.detach().clone() + trace.artifacts["stage2_audio_latent"] = stage2_audio.latent.detach().clone() + return LTX25ReferenceResult( + stage1_video, + stage1_audio, + stage2_video, + stage2_audio, + generator.get_state(), + trace, + ) + + def _encode_prompt(self, prompt: str) -> tuple[LTX25EmbeddingsProcessorOutput, dict[str, torch.Tensor]]: + hidden_states, token_ids, attention_mask = self.components.text_encoder.encode([prompt]) + artifacts: dict[str, torch.Tensor] = {} + if self.capture_prompt_intermediates: + processor = self.components.embeddings_processor + if isinstance(processor, _LazyCallable): + processor = processor.resolve() # type: ignore[assignment] + feature_extractor = getattr(processor, "feature_extractor", None) + if feature_extractor is not None: + video_features, audio_features = feature_extractor(hidden_states, attention_mask) + artifacts["video_features"] = video_features.detach().cpu() + artifacts["audio_features"] = audio_features.detach().cpu() + artifacts = { + "gemma_token_ids": token_ids.detach().cpu(), + "gemma_attention_mask": attention_mask.detach().cpu(), + **{ + f"gemma_hidden_state_{index}": hidden_state.detach().cpu() + for index, hidden_state in enumerate(hidden_states) + }, + **artifacts, + } + return self.components.embeddings_processor(hidden_states, attention_mask), artifacts + + def _tools( + self, request: LTX25ReferenceRequest, *, half_resolution: bool + ) -> tuple[VideoLatentTools, AudioLatentTools]: + height = request.height // 2 if half_resolution else request.height + width = request.width // 2 if half_resolution else request.width + pixel_shape = VideoPixelShape(1, request.num_frames, height, width, request.frame_rate) + video_shape = VideoLatentShape.from_pixel_shape(pixel_shape) + audio_shape = AudioLatentShape.from_duration(1, request.num_frames / request.frame_rate) + return ( + VideoLatentTools(VideoLatentPatchifier(patch_size=1), video_shape, request.frame_rate), + AudioLatentTools(AudioPatchifier(patch_size=1), audio_shape), + ) + + def _apply_image_conditions( + self, + state: LatentState, + tools: VideoLatentTools, + conditions: Sequence[LTX25ReferenceImageCondition], + height: int, + width: int, + ) -> LatentState: + if not conditions: + return state + encoder, release_after = self._video_encoder() + try: + for condition in conditions: + pixels = preprocess_ltx25_image( + condition.image, + height, + width, + self._image_crf(condition), + device=self.device, + dtype=self.dtype, + ) + encoded = encoder(pixels) + conditioning = ( + VideoConditionByLatentIndex(encoded, condition.strength, 0) + if condition.frame_idx == 0 + else VideoConditionByKeyframeIndex(encoded, condition.frame_idx, condition.strength) + ) + state = conditioning.apply_to(state, tools) + return state + finally: + if release_after and self.offload == "cpu": + encoder.to("cpu") # type: ignore[union-attr] + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def _image_crf(self, condition: LTX25ReferenceImageCondition) -> int: + if condition.crf is not None: + return condition.crf + return default_image_crf(self.video_encoder_path) + + def _video_encoder(self) -> tuple[VideoEncoder, bool]: + if self.components.video_encoder is not None: + return self.components.video_encoder, False + if self.video_encoder_path is None: + raise RuntimeError("I2V reference generation requires a video encoder or video_encoder_path") + from telefuser.models.ltx25.video_encoder import LTX25VideoEncoder + + return ( + LTX25VideoEncoder.from_checkpoint(self.video_encoder_path, device=self.device, torch_dtype=self.dtype), + True, + ) + + def _sample_stage( # noqa: PLR0913 + self, + *, + stage_name: str, + video: LatentState, + audio: LatentState, + video_context: torch.Tensor, + audio_context: torch.Tensor, + sigmas: torch.Tensor, + ancestral: bool, + seed: int, + trace: LTX25ReferenceTrace, + ) -> tuple[LatentState, LatentState]: + trace.artifacts[f"{stage_name}_initial_video_noise"] = video.latent.detach().clone() + trace.artifacts[f"{stage_name}_initial_audio_noise"] = audio.latent.detach().clone() + stepper = LTX25EulerAncestralStep() if ancestral else None + ancestral_generator = ancestral_noise_generator(seed, self.device) if ancestral else None + + for step_index in range(len(sigmas) - 1): + denoised_video, denoised_audio = self.components.transformer( + _modality(video, video_context, sigmas[step_index]), + _modality(audio, audio_context, sigmas[step_index]), + BatchedPerturbationConfig.empty(video.latent.shape[0]), + ) + if denoised_video is None or denoised_audio is None: + raise RuntimeError("LTX-2.5 joint T2V reference requires both video and audio x0 predictions") + denoised_video = _post_process(denoised_video, video) + denoised_audio = _post_process(denoised_audio, audio) + trace.artifacts[f"{stage_name}_step{step_index}_video_x0"] = denoised_video.detach().clone() + trace.artifacts[f"{stage_name}_step{step_index}_audio_x0"] = denoised_audio.detach().clone() + + if bool(sigmas[step_index + 1] == 0): + if stepper is None: + trace.artifacts[f"{stage_name}_step{step_index}_updated_video_latent"] = ( + denoised_video.detach().clone() + ) + trace.artifacts[f"{stage_name}_step{step_index}_updated_audio_latent"] = ( + denoised_audio.detach().clone() + ) + video = replace(video, latent=denoised_video.to(self.dtype)) + audio = replace(audio, latent=denoised_audio.to(self.dtype)) + continue + if stepper is None: + video_updated = _deterministic_euler_step(video.latent, denoised_video, sigmas, step_index) + audio_updated = _deterministic_euler_step(audio.latent, denoised_audio, sigmas, step_index) + trace.artifacts[f"{stage_name}_step{step_index}_updated_video_latent"] = video_updated.detach().clone() + trace.artifacts[f"{stage_name}_step{step_index}_updated_audio_latent"] = audio_updated.detach().clone() + video = replace(video, latent=video_updated) + audio = replace(audio, latent=audio_updated) + continue + + assert ancestral_generator is not None + video_noise = torch.randn( + video.latent.shape, + generator=ancestral_generator, + dtype=video.latent.dtype, + device=video.latent.device, + ) + audio_noise = torch.randn( + audio.latent.shape, + generator=ancestral_generator, + dtype=audio.latent.dtype, + device=audio.latent.device, + ) + trace.artifacts[f"{stage_name}_step{step_index}_ancestral_noise_video"] = video_noise.detach().clone() + trace.artifacts[f"{stage_name}_step{step_index}_ancestral_noise_audio"] = audio_noise.detach().clone() + video_updated = stepper.step(video.latent.float(), denoised_video, sigmas, step_index, video_noise) + audio_updated = stepper.step(audio.latent.float(), denoised_audio, sigmas, step_index, audio_noise) + trace.artifacts[f"{stage_name}_step{step_index}_updated_video_latent"] = video_updated.detach().clone() + trace.artifacts[f"{stage_name}_step{step_index}_updated_audio_latent"] = audio_updated.detach().clone() + video = replace(video, latent=_post_process(video_updated, video).to(self.dtype)) + audio = replace(audio, latent=_post_process(audio_updated, audio).to(self.dtype)) + return video, audio + + +__all__ = [ + "LTX25DistilledReference", + "LTX25ReferenceComponents", + "LTX25ReferenceImageCondition", + "LTX25ReferenceRequest", + "LTX25ReferenceResult", + "LTX25ReferenceTrace", +] diff --git a/telefuser/pipelines/ltx25_distilled/text_encoding.py b/telefuser/pipelines/ltx25_distilled/text_encoding.py new file mode 100644 index 0000000..e3916f8 --- /dev/null +++ b/telefuser/pipelines/ltx25_distilled/text_encoding.py @@ -0,0 +1,36 @@ +"""Prompt encoding and duration prediction for LTX-2.5.""" + +from __future__ import annotations + +import torch + +from telefuser.core.base_stage import BaseStage, with_model_offload +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.models.ltx25 import LTX25DurationHead, LTX25EmbeddingsProcessor, LTX25Gemma4TextEncoder +from telefuser.models.ltx25.duration import seconds_to_num_frames + + +class LTX25TextEncodingStage(BaseStage): + def __init__(self, module_manager: ModuleManager, config: ModelRuntimeConfig) -> None: + super().__init__("ltx25_text_encoding", config) + self.text_encoder: LTX25Gemma4TextEncoder = module_manager.fetch_module("ltx25_gemma4") + self.embeddings_processor: LTX25EmbeddingsProcessor = module_manager.fetch_module("ltx25_embeddings_processor") + self.duration_head: LTX25DurationHead = module_manager.fetch_module("ltx25_duration_head") + self.model_names = ["text_encoder", "embeddings_processor", "duration_head"] + + @with_model_offload(["text_encoder", "embeddings_processor", "duration_head"]) + @torch.inference_mode() + def encode(self, prompt: str) -> tuple[torch.Tensor, torch.Tensor]: + hidden_states, _, attention_mask = self.text_encoder.encode([prompt]) + encoded = self.embeddings_processor(hidden_states, attention_mask) + return encoded.video_encoding, encoded.audio_encoding + + @with_model_offload(["text_encoder", "embeddings_processor", "duration_head"]) + @torch.inference_mode() + def predict_num_frames(self, video_context: torch.Tensor, audio_context: torch.Tensor, frame_rate: float) -> int: + seconds = float(self.duration_head(video_context, audio_context).item()) + return seconds_to_num_frames(seconds, frame_rate=frame_rate) + + +__all__ = ["LTX25TextEncodingStage"] diff --git a/telefuser/pipelines/ltx25_distilled/video_conditioning.py b/telefuser/pipelines/ltx25_distilled/video_conditioning.py new file mode 100644 index 0000000..ab7c1c3 --- /dev/null +++ b/telefuser/pipelines/ltx25_distilled/video_conditioning.py @@ -0,0 +1,69 @@ +"""Still-image conditioning for LTX-2.5.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Protocol + +import torch +from PIL import Image + +from telefuser.core.base_stage import BaseStage, with_model_offload +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager + +from .image import default_image_crf, preprocess_ltx25_image +from .latent import LatentState, VideoConditionByKeyframeIndex, VideoConditionByLatentIndex, VideoLatentTools + + +class LTX25ImageConditionProtocol(Protocol): + image: Image.Image + frame_idx: int + strength: float + crf: int | None + + +class LTX25VideoConditioningStage(BaseStage): + def __init__(self, module_manager: ModuleManager, config: ModelRuntimeConfig) -> None: + super().__init__("ltx25_video_conditioning", config) + fetched = module_manager.fetch_module("ltx25_video_encoder", require_model_path=True) + if fetched is None: + raise ValueError("ModuleManager does not contain ltx25_video_encoder") + self.video_encoder, video_encoder_path = fetched + self.video_encoder_path = video_encoder_path + self.model_names = ["video_encoder"] + + @with_model_offload(["video_encoder"]) + @torch.inference_mode() + def apply( + self, + state: LatentState, + tools: VideoLatentTools, + conditions: Sequence[LTX25ImageConditionProtocol], + height: int, + width: int, + ) -> LatentState: + for condition in conditions: + if condition.frame_idx < 0: + raise ValueError(f"image frame_idx must be non-negative, got {condition.frame_idx}") + if not 0.0 <= condition.strength <= 1.0: + raise ValueError(f"image strength must be in [0, 1], got {condition.strength}") + pixels = preprocess_ltx25_image( + condition.image, + height, + width, + default_image_crf(self.video_encoder_path) if condition.crf is None else condition.crf, + device=self.device, + dtype=self.torch_dtype, + ) + encoded = self.video_encoder(pixels) + conditioning = ( + VideoConditionByLatentIndex(encoded, condition.strength, 0) + if condition.frame_idx == 0 + else VideoConditionByKeyframeIndex(encoded, condition.frame_idx, condition.strength) + ) + state = conditioning.apply_to(state, tools) + return state + + +__all__ = ["LTX25VideoConditioningStage"] diff --git a/telefuser/pipelines/ltx25_distilled/video_decoding.py b/telefuser/pipelines/ltx25_distilled/video_decoding.py new file mode 100644 index 0000000..1464513 --- /dev/null +++ b/telefuser/pipelines/ltx25_distilled/video_decoding.py @@ -0,0 +1,41 @@ +"""Video decoding for LTX-2.5.""" + +from __future__ import annotations + +from collections.abc import Iterable + +import torch + +from telefuser.core.base_stage import BaseStage, with_model_offload +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.models.ltx25 import DiffusionVideoDecoder, LTX25ConvVideoVAE + + +class LTX25VideoDecodingStage(BaseStage): + def __init__(self, module_manager: ModuleManager, config: ModelRuntimeConfig, *, video_vae: str) -> None: + super().__init__("ltx25_video_decoding", config) + self.video_decoder: DiffusionVideoDecoder | LTX25ConvVideoVAE = module_manager.fetch_module( + "ltx25_video_decoder" + ) + self.video_vae = video_vae + self.model_names = ["video_decoder"] + + @with_model_offload(["video_decoder"]) + @torch.inference_mode() + def decode(self, latent: torch.Tensor, generator: torch.Generator) -> tuple[torch.Tensor, ...]: + if self.video_vae == "diff": + return tuple(self.video_decoder.decode_video(latent, generator=generator)) # type: ignore[union-attr] + return _conv_video_chunks_to_rgb(self.video_decoder.decode(latent, generator=generator)) # type: ignore[union-attr] + + +def _conv_video_chunks_to_rgb(chunks: Iterable[torch.Tensor]) -> tuple[torch.Tensor, ...]: + output: list[torch.Tensor] = [] + for chunk in chunks: + if chunk.ndim != 5 or chunk.shape[0] != 1 or chunk.shape[1] != 3: + raise ValueError(f"LTX-2.5 ConvVAE decoder must return [1, 3, F, H, W], got {tuple(chunk.shape)}") + output.append(chunk[0].permute(1, 2, 3, 0).add(1).mul(0.5).clamp(0, 1)) + return tuple(output) + + +__all__ = ["LTX25VideoDecodingStage"] diff --git a/tests/assets/ltx25/README.md b/tests/assets/ltx25/README.md new file mode 100644 index 0000000..45dcbbd --- /dev/null +++ b/tests/assets/ltx25/README.md @@ -0,0 +1,12 @@ +# LTX-2.5 I2V Test Input + +`official_guitar_man.png` is the image directly referenced by the Lightricks +LTX-2 model card's image-to-video example: + +https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/guitar-man.png + +Its SHA-256 is +`e31cbbe4822ce07e1548121b436c0db3a067d1d78f2e75ab3e69375377b57274`. +The matching official example prompt is: `A man with short gray hair plays a +red electric guitar.` Formal LTX-2.5 I2V captures use this source at frame 0 +with strength 1.0. diff --git a/tests/assets/ltx25/official_guitar_man.png b/tests/assets/ltx25/official_guitar_man.png new file mode 100644 index 0000000..83ebe7d Binary files /dev/null and b/tests/assets/ltx25/official_guitar_man.png differ diff --git a/tests/unit/models/ltx25/test_audio.py b/tests/unit/models/ltx25/test_audio.py new file mode 100644 index 0000000..057e07e --- /dev/null +++ b/tests/unit/models/ltx25/test_audio.py @@ -0,0 +1,96 @@ +"""Tests for the isolated LTX-2.5 audio decoder and vocoder loader.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import torch +from safetensors.torch import save_file + +from telefuser.models.ltx25.audio import ( + AudioDecoderConfigurator, + VocoderConfigurator, + load_ltx25_audio_decoder_and_vocoder, + ltx25_audio_checkpoint_key_coverage, +) + + +def _audio_config() -> dict: + return { + "audio_vae": { + "model": { + "params": { + "ddconfig": { + "attn_resolutions": [], + "causality_axis": "height", + "ch": 32, + "ch_mult": [1, 2], + "mel_bins": 64, + "mid_block_add_attention": False, + "norm_type": "pixel", + "num_res_blocks": 1, + "out_ch": 2, + "resolution": 32, + "z_channels": 8, + }, + "sampling_rate": 16000, + } + }, + "preprocessing": {"stft": {"causal": True, "hop_length": 160}}, + }, + "vocoder": { + "vocoder": { + "activation": "snakebeta", + "resblock": "AMP1", + "resblock_dilation_sizes": [[1, 3, 5]], + "resblock_kernel_sizes": [3], + "upsample_initial_channel": 64, + "upsample_kernel_sizes": [4], + "upsample_rates": [2], + "use_bias_at_final": False, + "use_tanh_at_final": False, + }, + "bwe": { + "activation": "snakebeta", + "hop_length": 2, + "input_sampling_rate": 16000, + "n_fft": 8, + "num_mels": 4, + "output_sampling_rate": 48000, + "resblock": "AMP1", + "resblock_dilation_sizes": [[1, 3, 5]], + "resblock_kernel_sizes": [3], + "upsample_initial_channel": 32, + "upsample_kernel_sizes": [4], + "upsample_rates": [2], + "use_bias_at_final": False, + "use_tanh_at_final": False, + }, + }, + } + + +def test_audio_loader_strictly_maps_decoder_and_bwe_vocoder_weights(tmp_path: Path) -> None: + config = _audio_config() + decoder = AudioDecoderConfigurator.from_config(config) + vocoder = VocoderConfigurator.from_config(config) + state_dict = { + **{ + ("audio_vae.per_channel_statistics." if key.startswith("per_channel_statistics.") else "audio_vae.decoder.") + + key.removeprefix("per_channel_statistics."): value + for key, value in decoder.state_dict().items() + }, + **{f"vocoder.{key}": value for key, value in vocoder.state_dict().items()}, + } + checkpoint_path = tmp_path / "audio.safetensors" + save_file(state_dict, checkpoint_path, metadata={"config": json.dumps(config)}) + + loaded_decoder, loaded_vocoder = load_ltx25_audio_decoder_and_vocoder(checkpoint_path, torch_dtype=torch.float32) + coverage = ltx25_audio_checkpoint_key_coverage( + checkpoint_path, set(loaded_decoder.state_dict()), set(loaded_vocoder.state_dict()) + ) + assert coverage == (set(), set(), set(), set()) + assert loaded_vocoder.output_sampling_rate == 48000 + for expected, actual in zip(decoder.state_dict().values(), loaded_decoder.state_dict().values(), strict=True): + torch.testing.assert_close(actual, expected, equal_nan=True) diff --git a/tests/unit/models/ltx25/test_checkpoint.py b/tests/unit/models/ltx25/test_checkpoint.py new file mode 100644 index 0000000..25b4ed5 --- /dev/null +++ b/tests/unit/models/ltx25/test_checkpoint.py @@ -0,0 +1,68 @@ +"""LTX-2.5 split-checkpoint metadata contract tests.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch + +from telefuser.models.ltx25.checkpoint import ( + LTX25CheckpointMetadata as Metadata, +) +from telefuser.models.ltx25.checkpoint import ( + LTX25ModelPaths, + parse_model_version, + validate_gemma_source_checkpoint, +) +from telefuser.models.ltx25.gemma4 import _cast_checkpoint_tensor, gemma4_checkpoint_key_to_model_key + + +def test_model_version_parsing_normalizes_prerelease_separators() -> None: + assert parse_model_version("2.5") == (2, 5) + assert parse_model_version("2.5-rc1") == (2, 5) + assert parse_model_version("2.5.rc1") == (2, 5) + assert parse_model_version(None) == () + + +def test_split_layout_rejects_missing_required_checkpoint(tmp_path: Path) -> None: + with pytest.raises(FileNotFoundError, match="transformer"): + LTX25ModelPaths.from_model_root(tmp_path) + + +def test_gemma_source_checkpoint_requires_matching_gemma_version() -> None: + metadata = Metadata( + path=Path("transformer.safetensors"), + size_bytes=1, + sha256=None, + tensor_count=1, + metadata={"gemma_source_checkpoint": {"gemma_version": 4}}, + config={}, + model_version=(2, 5), + ) + validate_gemma_source_checkpoint(metadata, {"gemma_version": 4}) + with pytest.raises(ValueError, match="Gemma version mismatch"): + validate_gemma_source_checkpoint(metadata, {"gemma_version": 3}) + + +def test_gemma4_comfy_flat_checkpoint_keys_map_to_unified_model() -> None: + assert gemma4_checkpoint_key_to_model_key("model.layers.0.self_attn.q_proj.weight") == ( + "model.model.language_model.layers.0.self_attn.q_proj.weight" + ) + assert gemma4_checkpoint_key_to_model_key("model.embed_tokens.weight") == ( + "model.model.language_model.embed_tokens.weight" + ) + assert gemma4_checkpoint_key_to_model_key("audio_projector.embedding_projection.weight") == ( + "model.model.embed_audio.embedding_projection.weight" + ) + assert gemma4_checkpoint_key_to_model_key("tokenizer_json") is None + + +def test_gemma_checkpoint_cast_matches_upstream_builder_policy() -> None: + vector = torch.ones(2, dtype=torch.float32) + scalar = torch.ones((), dtype=torch.float32) + integer = torch.ones(2, dtype=torch.int64) + + assert _cast_checkpoint_tensor(vector, torch.bfloat16).dtype is torch.bfloat16 + assert _cast_checkpoint_tensor(scalar, torch.bfloat16).dtype is torch.float32 + assert _cast_checkpoint_tensor(integer, torch.bfloat16).dtype is torch.int64 diff --git a/tests/unit/models/ltx25/test_conv_video_vae.py b/tests/unit/models/ltx25/test_conv_video_vae.py new file mode 100644 index 0000000..a092ee1 --- /dev/null +++ b/tests/unit/models/ltx25/test_conv_video_vae.py @@ -0,0 +1,32 @@ +"""Tests for the isolated LTX-2.5 Conv VAE checkpoint mapping.""" + +from __future__ import annotations + +from pathlib import Path + +import torch +from safetensors.torch import save_file + +from telefuser.models.ltx25.conv_video_vae import ltx25_conv_video_vae_checkpoint_key_coverage + + +def test_conv_vae_mapping_duplicates_shared_statistics_for_encoder_and_decoder(tmp_path: Path) -> None: + checkpoint_path = tmp_path / "conv_video_vae.safetensors" + save_file( + { + "encoder.conv.weight": torch.ones(1), + "decoder.conv.weight": torch.ones(1), + "per_channel_statistics.mean-of-means": torch.zeros(1), + "per_channel_statistics.std-of-means": torch.ones(1), + }, + checkpoint_path, + ) + model_keys = { + "encoder.conv.weight", + "decoder.conv.weight", + "encoder.per_channel_statistics.mean-of-means", + "encoder.per_channel_statistics.std-of-means", + "decoder.per_channel_statistics.mean-of-means", + "decoder.per_channel_statistics.std-of-means", + } + assert ltx25_conv_video_vae_checkpoint_key_coverage(checkpoint_path, model_keys) == (set(), set()) diff --git a/tests/unit/models/ltx25/test_diffusion_video_decoder.py b/tests/unit/models/ltx25/test_diffusion_video_decoder.py new file mode 100644 index 0000000..7ecb86d --- /dev/null +++ b/tests/unit/models/ltx25/test_diffusion_video_decoder.py @@ -0,0 +1,193 @@ +"""Tests for isolated LTX-2.5 DiffVAE checkpoint mapping.""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace + +import torch +from safetensors.torch import save_file + +from telefuser.models.ltx25.diff_vae import diffusion_tiling +from telefuser.models.ltx25.diff_vae.diffusion_video_decoder import ( + DiffusionVideoDecoder, + _configure_chunked_compile_mode, + _configure_chunked_eager_mode, + ltx25_diffusion_vae_checkpoint_key_coverage, +) +from telefuser.models.ltx25.diff_vae.transformer import compiling as diffvae_compiling +from telefuser.models.ltx25.diff_vae.transformer.chunked.block import ChunkedDiffusionNABlock +from telefuser.models.ltx25.diff_vae.transformer.config import DiffVAEMode +from telefuser.ops import neighborhood_attention as attention_ops + + +def test_diffvae_mapping_splits_qkv_and_ignores_verified_unused_type_embedding(tmp_path: Path) -> None: + checkpoint_path = tmp_path / "diffusion_video_vae.safetensors" + save_file( + { + "decoder.block.qkv.weight": torch.ones(12, 4), + "decoder.block.qkv.bias": torch.ones(12), + "decoder.type_emb": torch.ones(4), + "per_channel_statistics.mean-of-means": torch.zeros(4), + "per_channel_statistics.std-of-means": torch.ones(4), + }, + checkpoint_path, + ) + model_keys = { + "block.qkv.to_q.weight", + "block.qkv.to_k.weight", + "block.qkv.to_v.weight", + "block.qkv.to_q.bias", + "block.qkv.to_k.bias", + "block.qkv.to_v.bias", + "per_channel_statistics.mean-of-means", + "per_channel_statistics.std-of-means", + } + assert ltx25_diffusion_vae_checkpoint_key_coverage(checkpoint_path, model_keys) == (set(), set()) + + +def test_chunked_eager_configuration_uses_deferred_stage4_and_width_chunks() -> None: + model = DiffusionVideoDecoder( + in_channels=8, + out_channels=3, + patch_size=1, + head_dim=8, + rope_dim_split=(2, 2, 4), + stage_channels=(8, 8, 8, 8, 8), + stage_depths=(1, 1, 1, 1, 1), + stage_kernels=((1, 1, 1),) * 5, + upsamples=(((1, 1, 1), 1),) * 4, + stage5_kernel=(1, 1, 1), + t_emb_dim=8, + ) + + _configure_chunked_eager_mode(model) + + assert model.deferred_stage4_upsample + for block in model.diff_blocks: + assert isinstance(block, ChunkedDiffusionNABlock) + assert block.stage4_upsample is model.upsamples[3] + assert block.attn.w_chunks == 4 + assert block.attn.rope_num_tiles == 1 + + +def test_chunked_compile_configuration_compiles_only_diffusion_residuals(monkeypatch) -> None: + compiled: list[object] = [] + kv_parallelism: list[bool] = [] + model = DiffusionVideoDecoder( + in_channels=8, + out_channels=3, + patch_size=1, + head_dim=8, + rope_dim_split=(2, 2, 4), + stage_channels=(8, 8, 8, 8, 8), + stage_depths=(1, 1, 1, 1, 1), + stage_kernels=((1, 1, 1),) * 5, + upsamples=(((1, 1, 1), 1),) * 4, + stage5_kernel=(1, 1, 1), + t_emb_dim=8, + ) + + def fake_compile(function, **kwargs): + compiled.append((function, kwargs)) + return function + + monkeypatch.setattr(diffvae_compiling.torch, "compile", fake_compile) + monkeypatch.setattr( + diffvae_compiling, + "configure_neighborhood_attention_kv_parallelism", + lambda enabled: kv_parallelism.append(enabled), + ) + + _configure_chunked_compile_mode(model) + + assert model.deferred_stage4_upsample + assert model.mark_dynamic_shapes + assert kv_parallelism == [False] + assert len(compiled) == len(model.diff_blocks) + for block in model.diff_blocks: + assert isinstance(block, ChunkedDiffusionNABlock) + assert block.stage4_upsample is model.upsamples[3] + assert block.attn.w_chunks == 4 + assert block.attn.natten_backend == "cutlass-fna" + + +def test_chunked_compile_tiling_uses_the_conservative_memory_budget() -> None: + assert diffusion_tiling.stage5_mem_coef(DiffVAEMode.CHUNKED_COMPILE) == 7 + assert diffusion_tiling.budget_safety_bytes(DiffVAEMode.CHUNKED_COMPILE) == 2 << 30 + + +def test_chunked_compile_tiling_caps_natten_workspace_tile_geometry() -> None: + tiling = diffusion_tiling.recommended_decode_tiling_config( + tile_halos=((0, 0, 0), (0, 0, 0)), + pixel_scale=diffusion_tiling.VIDEO_SCALE_FACTORS, + min_tile_size_s4=(1, 1, 1), + patch_size=1, + height=1536, + width=1536, + num_frames=121, + mode=DiffVAEMode.CHUNKED_COMPILE, + free_bytes=1 << 40, + stage5_channels=8, + stage4_channels=8, + upsample_strides=((1, 1, 1),) * 4, + ) + + assert tiling.frames.tile_size <= 80 + assert tiling.height.tile_size <= 320 + assert tiling.width.tile_size <= 320 + + +def test_natten_availability_requires_the_cuda_extension(monkeypatch) -> None: + monkeypatch.setattr(attention_ops, "_NATTEN_AVAILABLE", True) + monkeypatch.setattr(attention_ops, "natten", SimpleNamespace(HAS_LIBNATTEN=False)) + assert not attention_ops.natten_available() + + monkeypatch.setattr(attention_ops, "natten", SimpleNamespace(HAS_LIBNATTEN=True)) + assert attention_ops.natten_available() + + +def test_neighborhood_attention_dispatch_normalizes_qkv_dtype(monkeypatch) -> None: + captured: dict[str, object] = {} + + def fake_na3d(query, key, value, **kwargs): + captured.update(query=query, key=key, value=value, **kwargs) + return value + + monkeypatch.setattr(attention_ops, "_NATTEN_AVAILABLE", True) + monkeypatch.setattr(attention_ops, "natten", SimpleNamespace(HAS_LIBNATTEN=True, na3d=fake_na3d)) + + value = torch.zeros((1, 1, 1, 1, 1, 2), dtype=torch.bfloat16) + output = attention_ops.neighborhood_attention_3d( + torch.ones_like(value, dtype=torch.float32), + torch.ones_like(value, dtype=torch.float32), + value, + kernel_size=(1, 3, 3), + backend="cutlass-fna", + ) + + assert output is value + assert captured["query"].dtype == captured["key"].dtype == captured["value"].dtype == torch.bfloat16 + assert captured["kernel_size"] == (1, 3, 3) + + +def test_diffvae_recommends_auto_tiling_from_decoder_configuration() -> None: + model = DiffusionVideoDecoder( + in_channels=8, + out_channels=3, + patch_size=1, + head_dim=8, + rope_dim_split=(2, 2, 4), + stage_channels=(8, 8, 8, 8, 8), + stage_depths=(1, 1, 1, 1, 1), + stage_kernels=((1, 1, 1),) * 5, + upsamples=(((1, 1, 1), 1),) * 4, + stage5_kernel=(1, 1, 1), + t_emb_dim=8, + ) + + tiling = model.recommended_tiling_config(height=64, width=64, num_frames=9, free_bytes=1 << 40, model_bytes=0) + + assert tiling.frames.tile_size >= 9 + assert tiling.height.tile_size >= 64 + assert tiling.width.tile_size >= 64 diff --git a/tests/unit/models/ltx25/test_duration.py b/tests/unit/models/ltx25/test_duration.py new file mode 100644 index 0000000..dcae947 --- /dev/null +++ b/tests/unit/models/ltx25/test_duration.py @@ -0,0 +1,52 @@ +"""LTX-2.5 DurationHead loading and frame-grid contracts.""" + +from __future__ import annotations + +import os +from pathlib import Path + +import pytest +import torch + +from telefuser.models.ltx25.duration import ( + LTX25DurationHead, + ltx25_duration_checkpoint_key_coverage, + seconds_to_num_frames, +) + +_MODEL_ROOT = os.environ.get("LTX25_MODEL_ROOT") + + +def test_duration_head_accepts_either_connector_modality() -> None: + model = LTX25DurationHead( + video_cross_attention_dim=4, + audio_cross_attention_dim=2, + pooler_hidden_dim=4, + num_pooler_heads=2, + ) + + assert model(video_tokens=torch.zeros(1, 3, 4)).shape == (1,) + assert model(audio_tokens=torch.zeros(1, 3, 2)).shape == (1,) + with pytest.raises(ValueError, match="video_tokens or audio_tokens"): + model() + + +@pytest.mark.parametrize( + ("seconds", "frame_rate", "expected"), + [(0.2, 24.0, 25), (1.0, 24.0, 25), (1.5, 24.0, 33), (100.0, 24.0, 473)], +) +def test_duration_frame_resolution_matches_upstream_causal_grid( + seconds: float, frame_rate: float, expected: int +) -> None: + assert seconds_to_num_frames(seconds, frame_rate=frame_rate) == expected + + +@pytest.mark.skipif(_MODEL_ROOT is None, reason="LTX25_MODEL_ROOT is not configured") +def test_duration_checkpoint_has_full_strict_coverage() -> None: + checkpoint = Path(_MODEL_ROOT) / "model_patches/ltx-2.5-duration-head-bf16.safetensors" + model = LTX25DurationHead.from_checkpoint(checkpoint, device="cpu") + + unexpected, missing = ltx25_duration_checkpoint_key_coverage(checkpoint, set(model.state_dict())) + + assert not unexpected + assert not missing diff --git a/tests/unit/models/ltx25/test_spatial_upsampler.py b/tests/unit/models/ltx25/test_spatial_upsampler.py new file mode 100644 index 0000000..931eb50 --- /dev/null +++ b/tests/unit/models/ltx25/test_spatial_upsampler.py @@ -0,0 +1,54 @@ +"""Tests for the isolated LTX-2.5 spatial upsampler.""" + +from __future__ import annotations + +import json +from dataclasses import asdict +from pathlib import Path + +import torch +from safetensors.torch import save_file + +from telefuser.models.ltx25.spatial_upsampler import ( + LTX25PixelShuffleND, + LTX25SpatialUpsampler, + LTX25SpatialUpsamplerConfig, + load_video_latent_statistics, + upsample_video_latent, +) + + +def test_spatial_pixel_shuffle_matches_torch_2d_pixel_shuffle() -> None: + value = torch.arange(2 * 12 * 3 * 5, dtype=torch.float32).reshape(2, 12, 3, 5) + actual = LTX25PixelShuffleND(2)(value) + expected = torch.nn.functional.pixel_shuffle(value, upscale_factor=2) + torch.testing.assert_close(actual, expected) + + +def test_spatial_upsampler_loads_exact_checkpoint_and_statistics(tmp_path: Path) -> None: + config = LTX25SpatialUpsamplerConfig( + in_channels=4, + mid_channels=32, + num_blocks_per_stage=1, + dims=3, + spatial_upsample=True, + temporal_upsample=False, + ) + source = LTX25SpatialUpsampler(config) + upsampler_path = tmp_path / "upsampler.safetensors" + save_file(source.state_dict(), upsampler_path, metadata={"config": json.dumps(asdict(config))}) + loaded = LTX25SpatialUpsampler.from_checkpoint(upsampler_path, torch_dtype=torch.float32) + latent = torch.randn(1, 4, 2, 3, 5) + torch.testing.assert_close(loaded(latent), source(latent)) + + vae_path = tmp_path / "video_vae.safetensors" + save_file( + { + "per_channel_statistics.std-of-means": torch.full((4,), 2.0), + "per_channel_statistics.mean-of-means": torch.full((4,), 0.25), + }, + vae_path, + ) + statistics = load_video_latent_statistics(vae_path) + expected = statistics.normalize(source(statistics.un_normalize(latent))) + torch.testing.assert_close(upsample_video_latent(latent, source, statistics), expected) diff --git a/tests/unit/models/ltx25/test_transformer.py b/tests/unit/models/ltx25/test_transformer.py new file mode 100644 index 0000000..2dc05e2 --- /dev/null +++ b/tests/unit/models/ltx25/test_transformer.py @@ -0,0 +1,113 @@ +"""LTX-2.5 isolated AV-transformer architecture tests.""" + +from __future__ import annotations + +import torch + +from telefuser.core.config import AttentionConfig, AttnImplType +from telefuser.models.ltx25.transformer import ( + Attention, + LTX25AVTransformer, + Modality, + ltx25_transformer_key_to_model_key, +) + + +def _small_ltx25_config() -> dict: + return { + "transformer": { + "num_layers": 48, + "rope_type": "split", + "apply_gated_attention": True, + "ff_bias": False, + "caption_proj_before_connector": True, + "activation_fn": "gelu-approximate", + "attention_bias": True, + "num_vector_embeds": None, + "dropout": 0.0, + "num_embeds_ada_norm": 1000, + "use_linear_projection": False, + "only_cross_attention": False, + "cross_attention_norm": True, + "double_self_attention": False, + "upcast_attention": False, + "standardization_norm": "rms_norm", + "norm_elementwise_affine": False, + "qk_norm": "rms_norm", + "positional_embedding_type": "rope", + "use_audio_video_cross_attention": True, + "share_ff": False, + "av_cross_ada_norm": True, + "use_middle_indices_grid": True, + "num_attention_heads": 1, + "attention_head_dim": 4, + "in_channels": 4, + "out_channels": 4, + "cross_attention_dim": 4, + "audio_num_attention_heads": 1, + "audio_attention_head_dim": 4, + "audio_in_channels": 4, + "audio_out_channels": 4, + "audio_cross_attention_dim": 4, + "norm_eps": 1e-6, + "positional_embedding_theta": 10000.0, + "positional_embedding_max_pos": [20, 32, 32], + "audio_positional_embedding_max_pos": [20], + "timestep_scale_multiplier": 1000, + "av_ca_timestep_scale_multiplier": 1000.0, + "frequencies_precision": "float64", + "cross_attention_adaln": True, + "use_keyframes_abs_pos_embedding": True, + } + } + + +def test_ltx25_transformer_builds_metadata_architecture_with_asymmetric_ff_biases() -> None: + with torch.device("meta"): + model = LTX25AVTransformer(_small_ltx25_config()) + keys = set(model.state_dict()) + assert "velocity_model.transformer_blocks.0.ff.net.0.proj.bias" not in keys + assert "velocity_model.transformer_blocks.0.audio_ff.net.0.proj.bias" in keys + assert "velocity_model.keyframes_abs_pos_embedding" in keys + + +def test_transformer_split_checkpoint_mapping_excludes_embedding_processor_weights() -> None: + assert ( + ltx25_transformer_key_to_model_key("model.diffusion_model.proj_out.weight") == "velocity_model.proj_out.weight" + ) + assert ( + ltx25_transformer_key_to_model_key("model.diffusion_model.video_embeddings_connector.learnable_registers") + is None + ) + assert ltx25_transformer_key_to_model_key("unrelated.weight") is None + + +def test_ltx25_transformer_exposes_the_runtime_attention_override() -> None: + with torch.device("meta"): + model = LTX25AVTransformer(_small_ltx25_config()) + + model.set_attention_config(AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA)) + assert Attention.attention_config.attn_impl is AttnImplType.TORCH_SDPA + + +def test_cross_attention_scale_shift_uses_per_token_timesteps() -> None: + model = LTX25AVTransformer(_small_ltx25_config()) + video = Modality( + latent=torch.randn(1, 3, 4), + sigma=torch.tensor([1.0]), + timesteps=torch.ones(1, 3), + positions=torch.zeros(1, 3, 3, 2), + context=torch.randn(1, 2, 4), + ) + audio = Modality( + latent=torch.randn(1, 2, 4), + sigma=torch.tensor([1.0]), + timesteps=torch.ones(1, 2), + positions=torch.zeros(1, 1, 2, 2), + context=torch.randn(1, 2, 4), + ) + + args = model.velocity_model.audio_args_preprocessor.prepare(audio, video) + + assert args.cross_scale_shift_timestep is not None + assert args.cross_scale_shift_timestep.shape == (1, 2, 16) diff --git a/tests/unit/models/ltx25/test_video_encoder.py b/tests/unit/models/ltx25/test_video_encoder.py new file mode 100644 index 0000000..ce211c0 --- /dev/null +++ b/tests/unit/models/ltx25/test_video_encoder.py @@ -0,0 +1,44 @@ +"""Tests for isolated LTX-2.5 video encoder checkpoint mapping.""" + +from __future__ import annotations + +from pathlib import Path + +import torch +from safetensors.torch import save_file + +from telefuser.models.ltx25.video_encoder import _video_encoder_kwargs, ltx25_video_encoder_checkpoint_key_coverage + + +def test_video_encoder_mapping_accepts_encoder_and_shared_statistics(tmp_path: Path) -> None: + checkpoint_path = tmp_path / "video_vae.safetensors" + save_file( + { + "encoder.conv_in.weight": torch.ones(1), + "per_channel_statistics.mean-of-means": torch.zeros(1), + "per_channel_statistics.std-of-means": torch.ones(1), + }, + checkpoint_path, + ) + model_keys = { + "conv_in.weight", + "per_channel_statistics.mean-of-means", + "per_channel_statistics.std-of-means", + } + assert ltx25_video_encoder_checkpoint_key_coverage(checkpoint_path, model_keys) == (set(), set()) + + +def test_video_encoder_uses_nested_vae_latent_channels() -> None: + kwargs = _video_encoder_kwargs( + { + "vae": { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 128, + "encoder_blocks": [["res_x", {"num_layers": 4}]], + } + } + ) + + assert kwargs["in_channels"] == 3 + assert kwargs["out_channels"] == 128 diff --git a/tests/unit/models/test_ltx25_sequence_parallel.py b/tests/unit/models/test_ltx25_sequence_parallel.py new file mode 100644 index 0000000..efb012f --- /dev/null +++ b/tests/unit/models/test_ltx25_sequence_parallel.py @@ -0,0 +1,121 @@ +"""Sequence-parallel contracts for the isolated LTX-2.5 transformer.""" + +from __future__ import annotations + +import torch + +from telefuser.core.config import AttnImplType +from telefuser.models.ltx25.transformer import Attention, LTXModel, LTXRopeType, TransformerArgs + + +def _tiny_model() -> LTXModel: + return LTXModel( + num_attention_heads=4, + attention_head_dim=8, + in_channels=8, + out_channels=8, + num_layers=1, + cross_attention_dim=16, + audio_num_attention_heads=4, + audio_attention_head_dim=8, + audio_in_channels=8, + audio_out_channels=8, + audio_cross_attention_dim=16, + ) + + +def test_enable_usp_wires_self_and_av_attention_only(monkeypatch) -> None: + model = _tiny_model() + mesh = object() + group = object() + monkeypatch.setattr("telefuser.models.ltx25.transformer.get_attention_strategy", lambda _: "ulysses") + monkeypatch.setattr("telefuser.models.ltx25.transformer.get_ulysses_world_size", lambda _: 4) + monkeypatch.setattr("telefuser.models.ltx25.transformer.get_ulysses_group", lambda _: group) + + model.enable_usp(mesh) # type: ignore[arg-type] + + block = model.transformer_blocks[0] + assert model.usp_flag + assert block.attn1.ulysses_group is group + assert block.audio_attn1.ulysses_group is group + assert block.audio_to_video_attn.ulysses_group is group + assert block.video_to_audio_attn.ulysses_group is group + assert block.attn2.ulysses_group is None + assert block.audio_attn2.ulysses_group is None + + +def test_shard_transformer_args_masks_sequence_padding(monkeypatch) -> None: + model = _tiny_model() + model.device_mesh = object() # type: ignore[assignment] + monkeypatch.setattr("telefuser.models.ltx25.transformer.get_ulysses_world_size", lambda _: 4) + shard_calls = [] + monkeypatch.setattr( + "telefuser.models.ltx25.transformer.sequence_parallel_shard", + lambda mesh, tensors, dimensions: shard_calls.append((mesh, tensors, dimensions)), + ) + args = TransformerArgs( + x=torch.zeros(1, 5, 32), + context=torch.zeros(1, 2, 16), + context_mask=torch.zeros(1, 1, 1, 2), + timesteps=torch.zeros(1, 5, 6, 32), + embedded_timestep=torch.zeros(1, 5, 32), + positional_embeddings=(torch.zeros(1, 5, 32), torch.zeros(1, 5, 32)), + cross_positional_embeddings=(torch.zeros(1, 5, 16), torch.zeros(1, 5, 16)), + cross_scale_shift_timestep=torch.zeros(1, 5, 4, 32), + cross_gate_timestep=torch.zeros(1, 1, 32), + enabled=True, + ) + + sharded, sequence_length = model._shard_transformer_args(args) + + assert sequence_length == 5 + assert shard_calls[0][2] == [1] * 8 + assert sharded.key_padding_mask.shape == (1, 1, 1, 8) + assert torch.all(sharded.key_padding_mask[..., :5] == 0) + assert torch.all(sharded.key_padding_mask[..., 5:] == torch.finfo(torch.float32).min) + assert sharded.self_attention_mask is sharded.key_padding_mask + + +def test_split_rope_shards_the_token_dimension_when_heads_match_sequence_length(monkeypatch) -> None: + model = _tiny_model() + model.rope_type = LTXRopeType.SPLIT + model.device_mesh = object() # type: ignore[assignment] + monkeypatch.setattr("telefuser.models.ltx25.transformer.get_ulysses_world_size", lambda _: 2) + shard_calls = [] + monkeypatch.setattr( + "telefuser.models.ltx25.transformer.sequence_parallel_shard", + lambda mesh, tensors, dimensions: shard_calls.append((mesh, tensors, dimensions)), + ) + args = TransformerArgs( + x=torch.zeros(1, 4, 32), + context=torch.zeros(1, 2, 16), + context_mask=torch.zeros(1, 1, 1, 2), + timesteps=torch.zeros(1, 4, 6, 32), + embedded_timestep=torch.zeros(1, 4, 32), + positional_embeddings=(torch.zeros(1, 4, 4, 4), torch.zeros(1, 4, 4, 4)), + cross_positional_embeddings=(torch.zeros(1, 4, 4, 4), torch.zeros(1, 4, 4, 4)), + cross_scale_shift_timestep=torch.zeros(1, 4, 4, 32), + cross_gate_timestep=torch.zeros(1, 1, 32), + enabled=True, + ) + + model._shard_transformer_args(args) + + assert shard_calls[0][2] == [1, 1, 1, 1, 2, 2, 2, 2] + + +def test_padding_mask_forces_mask_aware_attention(monkeypatch) -> None: + selected_backends = [] + + def fake_attention(query, key, value, *, attention_config, **kwargs): + del key, value, kwargs + selected_backends.append(attention_config.attn_impl) + return query + + monkeypatch.setattr("telefuser.models.ltx25.transformer.attn_func", fake_attention) + attention = Attention(query_dim=8, heads=2, dim_head=4) + + output = attention(torch.zeros(1, 4, 8), mask=torch.zeros(1, 1, 1, 4), enforce_mask=True) + + assert output.shape == (1, 4, 8) + assert selected_backends == [AttnImplType.TORCH_SDPA] diff --git a/tests/unit/pipelines/ltx25_distilled/__init__.py b/tests/unit/pipelines/ltx25_distilled/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/pipelines/ltx25_distilled/test_artifact_comparison.py b/tests/unit/pipelines/ltx25_distilled/test_artifact_comparison.py new file mode 100644 index 0000000..a0bd7c3 --- /dev/null +++ b/tests/unit/pipelines/ltx25_distilled/test_artifact_comparison.py @@ -0,0 +1,111 @@ +"""LTX-2.5 artifact-comparison contract tests.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import torch + +from tools.validation.compare_ltx25_artifacts import compare_captures + + +def _write_capture( + root: Path, + artifacts: dict[str, torch.Tensor], + *, + checkpoints: dict[str, object] | None = None, + audio: dict[str, object] | None = None, +) -> None: + root.mkdir(exist_ok=True) + descriptors = {} + for name, value in artifacts.items(): + path = root / f"{name}.pt" + torch.save(value, path) + descriptors[name] = {"path": path.name} + (root / "capture_manifest.json").write_text( + json.dumps( + { + "request": {"seed": 7}, + "checkpoints": checkpoints or {}, + "audio": audio or {}, + "artifacts": descriptors, + } + ), + encoding="utf-8", + ) + + +def test_comparison_requires_exact_noise_and_accepts_close_float_artifacts(tmp_path: Path) -> None: + golden = tmp_path / "golden" + candidate = tmp_path / "candidate" + _write_capture(golden, {"stage1_initial_video_noise": torch.ones(2), "latent": torch.ones(2)}) + _write_capture(candidate, {"stage1_initial_video_noise": torch.ones(2), "latent": torch.tensor([1.0, 1.0001])}) + report = compare_captures(golden, candidate) + assert report["passed"] + + _write_capture(candidate, {"stage1_initial_video_noise": torch.tensor([1.0, 0.0]), "latent": torch.ones(2)}) + report = compare_captures(golden, candidate) + assert not report["passed"] + assert "stage1_initial_video_noise" in report["failures"] + + +def test_comparison_accepts_exact_non_contract_tensors_without_metric_margin(tmp_path: Path) -> None: + golden = tmp_path / "golden" + candidate = tmp_path / "candidate" + _write_capture(golden, {"decoded_rgb": torch.full((1024,), 0.5, dtype=torch.bfloat16)}) + _write_capture(candidate, {"decoded_rgb": torch.full((1024,), 0.5, dtype=torch.bfloat16)}) + + report = compare_captures(golden, candidate, cosine_threshold=1.1) + + assert report["passed"] + assert report["tensors"]["decoded_rgb"]["exact"] + assert report["tensors"]["decoded_rgb"]["ssim"] == 1.0 + + +def test_comparison_requires_matching_checkpoint_and_audio_contracts(tmp_path: Path) -> None: + """Golden checkpoint and decoded-audio metadata are exact capture contracts.""" + golden = tmp_path / "golden" + candidate = tmp_path / "candidate" + artifacts = {"decoded_waveform": torch.ones(2, 32)} + _write_capture( + golden, + artifacts, + checkpoints={"transformer": {"sha256": "golden"}}, + audio={"sample_rate": 48000}, + ) + _write_capture( + candidate, + artifacts, + checkpoints={"transformer": {"sha256": "candidate"}}, + audio={"sample_rate": 44100}, + ) + + report = compare_captures(golden, candidate) + + assert not report["passed"] + assert not report["checkpoint_match"] + assert not report["audio_match"] + assert "checkpoints" in report["failures"] + assert "audio" in report["failures"] + + +def test_comparison_allows_upstream_only_prompt_diagnostics(tmp_path: Path) -> None: + golden = tmp_path / "golden" + candidate = tmp_path / "candidate" + _write_capture( + golden, + { + "latent": torch.ones(2), + "gemma_token_ids": torch.ones(2, dtype=torch.long), + "gemma_hidden_state_0": torch.ones(2), + }, + ) + _write_capture(candidate, {"latent": torch.ones(2)}) + + report = compare_captures(golden, candidate) + + assert report["passed"] + assert not report["artifact_set_match"] + assert report["artifact_contract_match"] + assert not report["missing_required_from_candidate"] diff --git a/tests/unit/pipelines/ltx25_distilled/test_core.py b/tests/unit/pipelines/ltx25_distilled/test_core.py new file mode 100644 index 0000000..bd76a83 --- /dev/null +++ b/tests/unit/pipelines/ltx25_distilled/test_core.py @@ -0,0 +1,148 @@ +"""Unit tests for isolated LTX-2.5 distilled sampling helpers.""" + +from __future__ import annotations + +from dataclasses import replace + +import torch + +from telefuser.models.ltx25.sampler import LTX25EulerAncestralStep +from telefuser.pipelines.ltx25_distilled.core import ( + LTX25SimpleDenoiser, + euler_ancestral_denoising_loop, + euler_denoising_loop, + modality_from_latent_state, +) +from telefuser.pipelines.ltx25_distilled.latent import LatentState + + +class _IdentityDenoiserModel: + def __call__( + self, video: object, audio: object, perturbations: object + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + del perturbations + return ( + video.latent if video is not None else None, # type: ignore[union-attr] + audio.latent if audio is not None else None, # type: ignore[union-attr] + ) + + +class _ConstantDenoiserModel: + def __call__( + self, video: object, audio: object, perturbations: object + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + del perturbations + return ( + torch.ones_like(video.latent) if video is not None else None, # type: ignore[union-attr] + torch.ones_like(audio.latent) if audio is not None else None, # type: ignore[union-attr] + ) + + +def _state(latent: torch.Tensor, *, denoise_mask: torch.Tensor | None = None) -> LatentState: + mask = denoise_mask if denoise_mask is not None else torch.ones_like(latent[..., :1]) + return LatentState( + latent=latent, + denoise_mask=mask, + positions=torch.zeros(latent.shape[0], 1, latent.shape[1], 2), + clean_latent=torch.full_like(latent, 9.0), + ) + + +def test_modality_uses_masked_per_token_timesteps() -> None: + state = _state(torch.zeros(1, 2, 3), denoise_mask=torch.tensor([[[1.0], [0.0]]])) + + modality = modality_from_latent_state(state, torch.zeros(1, 4, 5), torch.tensor([0.75])) + + torch.testing.assert_close(modality.timesteps, torch.tensor([[[0.75], [0.0]]])) + assert modality.positions is state.positions + + +def test_euler_loop_preserves_conditioning_tokens() -> None: + state = _state(torch.zeros(1, 2, 1), denoise_mask=torch.tensor([[[1.0], [0.0]]])) + denoiser = LTX25SimpleDenoiser(torch.zeros(1, 1, 1), None) + + result, audio = euler_denoising_loop( + torch.tensor([0.5, 0.0]), + state, + None, + _ConstantDenoiserModel(), # type: ignore[arg-type] + denoiser, + model_dtype=torch.float32, + ) + + assert audio is None + assert result is not None + torch.testing.assert_close(result.latent, torch.tensor([[[1.0], [9.0]]])) + + +def test_euler_loop_matches_upstream_scalar_sigma_velocity_rounding() -> None: + torch.manual_seed(23) + state = _state(torch.randn(1, 256, 1, dtype=torch.bfloat16)) + sigmas = torch.tensor([0.725, 0.421875]) + denoiser = LTX25SimpleDenoiser(torch.zeros(1, 1, 1), None) + + result, _ = euler_denoising_loop( + sigmas, + state, + None, + _ConstantDenoiserModel(), # type: ignore[arg-type] + denoiser, + model_dtype=torch.bfloat16, + ) + expected_velocity = ((state.latent.float() - torch.ones_like(state.latent).float()) / sigmas[0].item()).to( + torch.bfloat16 + ) + expected = (state.latent.float() + expected_velocity.float() * (sigmas[1] - sigmas[0])).to(torch.bfloat16) + + assert result is not None + torch.testing.assert_close(result.latent, expected) + + +def test_noising_preserves_clean_conditioning_tokens() -> None: + from telefuser.pipelines.ltx25_distilled.pipeline import _noised_state + + state = _state(torch.tensor([[[2.0], [0.0]]]), denoise_mask=torch.tensor([[[0.0], [1.0]]])) + noised = _noised_state(state, 1.0, torch.Generator().manual_seed(17)) + + torch.testing.assert_close(noised.latent[:, :1], state.clean_latent[:, :1]) + assert not torch.equal(noised.latent[:, 1:], state.clean_latent[:, 1:]) + + +def test_ancestral_loop_draws_video_noise_before_audio_noise() -> None: + video_state = _state(torch.zeros(1, 2, 1)) + audio_state = _state(torch.zeros(1, 3, 1)) + sigmas = torch.tensor([0.8, 0.4, 0.0]) + seed = 123 + denoiser = LTX25SimpleDenoiser(torch.zeros(1, 1, 1), torch.zeros(1, 1, 1)) + + video_result, audio_result = euler_ancestral_denoising_loop( + sigmas, + video_state, + audio_state, + _IdentityDenoiserModel(), # type: ignore[arg-type] + denoiser, + noise_seed=seed, + model_dtype=torch.float32, + ) + + generator = torch.Generator().manual_seed(seed) + stepper = LTX25EulerAncestralStep() + expected_video = stepper.step( + video_state.latent.float(), + video_state.latent, + sigmas, + 0, + torch.randn(video_state.latent.shape, generator=generator), + ) + expected_audio = stepper.step( + audio_state.latent.float(), + audio_state.latent, + sigmas, + 0, + torch.randn(audio_state.latent.shape, generator=generator), + ) + + assert video_result is not None + assert audio_result is not None + torch.testing.assert_close(video_result.latent, expected_video) + torch.testing.assert_close(audio_result.latent, expected_audio) diff --git a/tests/unit/pipelines/ltx25_distilled/test_image.py b/tests/unit/pipelines/ltx25_distilled/test_image.py new file mode 100644 index 0000000..2b0b9ea --- /dev/null +++ b/tests/unit/pipelines/ltx25_distilled/test_image.py @@ -0,0 +1,28 @@ +"""LTX-2.5 image-conditioning preprocessing tests.""" + +from __future__ import annotations + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + +from telefuser.pipelines.ltx25_distilled.image import preprocess_ltx25_image + + +def test_preprocess_ltx25_image_preserves_aspect_before_center_crop() -> None: + pixels = np.array( + [ + [[0, 0, 0], [64, 64, 64]], + [[128, 128, 128], [255, 255, 255]], + ], + dtype=np.uint8, + ) + image = Image.fromarray(pixels) + + actual = preprocess_ltx25_image(image, height=2, width=4, crf=0, dtype=torch.float32) + expected = torch.from_numpy(pixels).permute(2, 0, 1).unsqueeze(0).float() + expected = F.interpolate(expected, size=(4, 4), mode="bilinear", align_corners=False) + expected = (expected[:, :, 1:3] / 127.5 - 1.0).unsqueeze(2) + + torch.testing.assert_close(actual, expected) diff --git a/tests/unit/pipelines/ltx25_distilled/test_latent.py b/tests/unit/pipelines/ltx25_distilled/test_latent.py new file mode 100644 index 0000000..12b1a04 --- /dev/null +++ b/tests/unit/pipelines/ltx25_distilled/test_latent.py @@ -0,0 +1,48 @@ +"""LTX-2.5 latent-state construction contracts.""" + +from __future__ import annotations + +import torch + +from telefuser.models.ltx25.diff_vae.types import VideoLatentShape +from telefuser.pipelines.ltx25_distilled.latent import ( + AudioLatentShape, + AudioLatentTools, + AudioPatchifier, + VideoLatentPatchifier, + VideoLatentTools, +) + + +def test_video_state_keeps_upstream_float32_pixel_positions() -> None: + tools = VideoLatentTools( + patchifier=VideoLatentPatchifier(patch_size=1), + target_shape=VideoLatentShape(batch=1, channels=128, frames=2, height=1, width=1), + fps=24.0, + ) + + state = tools.create_initial_state(device=torch.device("cpu"), dtype=torch.bfloat16) + + assert state.positions.dtype == torch.float32 + assert state.keyframes_mask is not None + torch.testing.assert_close(state.keyframes_mask, torch.tensor([[[1.0], [0.0]]])) + torch.testing.assert_close( + state.positions[:, 0, :, :], + torch.tensor([[[0.0, 1.0 / 24.0], [1.0 / 24.0, 9.0 / 24.0]]]), + ) + torch.testing.assert_close( + state.positions[:, 1:, :, :], + torch.tensor([[[[0.0, 32.0], [0.0, 32.0]], [[0.0, 32.0], [0.0, 32.0]]]]), + ) + + +def test_audio_state_keeps_upstream_float32_time_positions() -> None: + tools = AudioLatentTools( + patchifier=AudioPatchifier(patch_size=1), + target_shape=AudioLatentShape(batch=1, channels=8, frames=2, mel_bins=16), + ) + + state = tools.create_initial_state(device=torch.device("cpu"), dtype=torch.bfloat16) + + assert state.positions.dtype == torch.float32 + torch.testing.assert_close(state.positions, torch.tensor([[[[0.0, 0.01], [0.01, 0.05]]]])) diff --git a/tests/unit/pipelines/ltx25_distilled/test_loader.py b/tests/unit/pipelines/ltx25_distilled/test_loader.py new file mode 100644 index 0000000..2d37696 --- /dev/null +++ b/tests/unit/pipelines/ltx25_distilled/test_loader.py @@ -0,0 +1,72 @@ +"""LTX-2.5 model-pack loading contracts.""" + +from __future__ import annotations + +from pathlib import Path + +import torch + +from telefuser.core.config import WeightOffloadType +from telefuser.core.module_manager import ModuleManager +from telefuser.models.ltx25 import LTX25ModelPaths +from telefuser.pipelines.ltx25_distilled import loader +from telefuser.pipelines.ltx25_distilled.pipeline import build_ltx25_distilled_config + + +class _Module(torch.nn.Module): + pass + + +def test_model_pack_loader_registers_every_component_with_module_manager(monkeypatch) -> None: + paths = LTX25ModelPaths( + *(Path(name) for name in ("transformer", "text", "diff", "conv", "audio", "up", "duration")) + ) + monkeypatch.setattr(LTX25ModelPaths, "from_model_root", staticmethod(lambda model_root: paths)) + + single_checkpoint_classes = ( + loader.LTX25Gemma4TextEncoder, + loader.LTX25DurationHead, + loader.LTX25VideoEncoder, + loader.LTX25AVTransformer, + loader.LTX25SpatialUpsampler, + loader.DiffusionVideoDecoder, + ) + for model_class in single_checkpoint_classes: + monkeypatch.setattr(model_class, "from_checkpoint", staticmethod(lambda *args, **kwargs: _Module())) + monkeypatch.setattr( + loader.LTX25EmbeddingsProcessor, + "from_checkpoints", + staticmethod(lambda *args, **kwargs: _Module()), + ) + monkeypatch.setattr(loader, "load_video_latent_statistics", lambda path: _Module()) + monkeypatch.setattr( + loader, + "load_ltx25_audio_decoder_and_vocoder", + lambda *args, **kwargs: (_Module(), _Module()), + ) + + manager = ModuleManager(device="cpu", torch_dtype=torch.bfloat16) + actual_paths = loader.load_ltx25_distilled_modules(manager, "unused", video_vae="diff") + + assert actual_paths is paths + assert manager.module_names == [ + "ltx25_gemma4", + "ltx25_embeddings_processor", + "ltx25_duration_head", + "ltx25_video_encoder", + "ltx25_transformer", + "ltx25_spatial_upsampler", + "ltx25_video_latent_statistics", + "ltx25_video_decoder", + "ltx25_audio_decoder", + "ltx25_vocoder", + ] + assert all(isinstance(module, _Module) for module in manager.modules) + + +def test_cpu_offload_config_uses_async_denoiser_and_model_offload_for_other_stages() -> None: + config = build_ltx25_distilled_config("cuda", torch.bfloat16, "diff", "cpu") + + assert config.denoising_config.offload_config.offload_type == WeightOffloadType.ASYNC_CPU_OFFLOAD + assert config.text_encoding_config.offload_config.offload_type == WeightOffloadType.MODEL_CPU_OFFLOAD + assert config.video_decoding_config.offload_config.offload_type == WeightOffloadType.MODEL_CPU_OFFLOAD diff --git a/tests/unit/pipelines/ltx25_distilled/test_pipeline.py b/tests/unit/pipelines/ltx25_distilled/test_pipeline.py new file mode 100644 index 0000000..461bd01 --- /dev/null +++ b/tests/unit/pipelines/ltx25_distilled/test_pipeline.py @@ -0,0 +1,209 @@ +"""ModuleManager and stage-composition contracts for LTX-2.5.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from examples.ltx25_distilled import ltx25_distilled_i2v_h100 as i2v_example +from telefuser.core.config import AttnImplType, WeightOffloadType +from telefuser.core.module_manager import ModuleManager +from telefuser.pipelines.ltx25_distilled.pipeline import ( + LTX25DistilledPipeline, + build_ltx25_distilled_config, +) + + +class _TextEncoder(torch.nn.Module): + def encode(self, prompts: list[str]) -> tuple[tuple[torch.Tensor, ...], torch.Tensor, torch.Tensor]: + batch = len(prompts) + hidden = torch.ones(batch, 2, 4) + return (hidden,), torch.ones(batch, 2, dtype=torch.long), torch.ones(batch, 2, dtype=torch.long) + + +class _Embeddings(torch.nn.Module): + def forward(self, hidden: tuple[torch.Tensor, ...], mask: torch.Tensor) -> SimpleNamespace: + del mask + return SimpleNamespace(video_encoding=hidden[0], audio_encoding=hidden[0]) + + +class _DurationHead(torch.nn.Module): + def forward(self, video: torch.Tensor, audio: torch.Tensor) -> torch.Tensor: + del video, audio + return torch.tensor([1.0]) + + +class _Transformer(torch.nn.Module): + def set_attention_config(self, attention_config: object) -> None: + self.attention_config = attention_config + + def forward( + self, video: object, audio: object, perturbations: object + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + del perturbations + return ( + torch.zeros_like(video.latent) if video is not None else None, # type: ignore[union-attr] + torch.zeros_like(audio.latent) if audio is not None else None, # type: ignore[union-attr] + ) + + +class _Upsampler(torch.nn.Module): + def forward(self, latent: torch.Tensor) -> torch.Tensor: + return latent.repeat_interleave(2, dim=3).repeat_interleave(2, dim=4) + + +class _Statistics(torch.nn.Module): + def un_normalize(self, latent: torch.Tensor) -> torch.Tensor: + return latent + + def normalize(self, latent: torch.Tensor) -> torch.Tensor: + return latent + + +class _VideoDecoder(torch.nn.Module): + def decode_video(self, latent: torch.Tensor, generator: torch.Generator) -> tuple[torch.Tensor, ...]: + del generator + return (latent,) + + +class _Identity(torch.nn.Module): + def forward(self, value: torch.Tensor) -> torch.Tensor: + return value + + +def _module_manager() -> ModuleManager: + manager = ModuleManager(device="cpu", torch_dtype=torch.bfloat16) + modules = { + "ltx25_gemma4": _TextEncoder(), + "ltx25_embeddings_processor": _Embeddings(), + "ltx25_duration_head": _DurationHead(), + "ltx25_video_encoder": _Identity(), + "ltx25_transformer": _Transformer(), + "ltx25_spatial_upsampler": _Upsampler(), + "ltx25_video_latent_statistics": _Statistics(), + "ltx25_video_decoder": _VideoDecoder(), + "ltx25_audio_decoder": _Identity(), + "ltx25_vocoder": _Identity(), + } + for name, module in modules.items(): + manager.add_module(module, name, path="unused") + return manager + + +def test_pipeline_composes_six_manager_backed_stages() -> None: + pipeline = LTX25DistilledPipeline(device="cpu", torch_dtype=torch.bfloat16) + pipeline.init( + _module_manager(), + build_ltx25_distilled_config("cpu", torch.bfloat16, "diff", "none"), + ) + + assert [stage.name for stage in pipeline._get_stages()] == [ + "ltx25_text_encoding", + "ltx25_video_conditioning", + "ltx25_denoising", + "ltx25_latent_upsampling", + "ltx25_video_decoding", + "ltx25_audio_decoding", + ] + assert len(pipeline._model_info) == 10 + + +def test_pipeline_runs_two_stage_contract_with_manager_owned_modules() -> None: + pipeline = LTX25DistilledPipeline(device="cpu", torch_dtype=torch.bfloat16) + pipeline.init( + _module_manager(), + build_ltx25_distilled_config("cpu", torch.bfloat16, "diff", "none"), + ) + + result = pipeline( + "A test prompt", + seed=7, + height=256, + width=384, + num_frames=9, + frame_rate=24.0, + ) + + assert result.video_latent.shape == (1, 128, 2, 8, 12) + assert result.audio_latent.shape == (1, 8, 9, 16) + assert result.video_chunks == (result.video_latent,) + assert result.audio.shape == result.audio_latent.squeeze(0).shape + assert result.audio.dtype == torch.float32 + + +def test_pipeline_resolves_auto_duration_through_text_stage() -> None: + pipeline = LTX25DistilledPipeline(device="cpu", torch_dtype=torch.bfloat16) + pipeline.init( + _module_manager(), + build_ltx25_distilled_config("cpu", torch.bfloat16, "diff", "none"), + ) + + result = pipeline("A test prompt", seed=7, height=256, width=384) + + assert result.num_frames == 25 + assert result.video_latent.shape == (1, 128, 4, 8, 12) + + +def test_build_config_supports_ulysses_and_attention_selection() -> None: + config = build_ltx25_distilled_config( + "cuda", + torch.bfloat16, + "diff", + "cpu", + parallelism=4, + attn_impl=AttnImplType.TORCH_SDPA, + ) + + denoising = config.denoising_config + assert denoising.attention_config.attn_impl == AttnImplType.TORCH_SDPA + assert denoising.parallel_config.device_ids == [0, 1, 2, 3] + assert denoising.parallel_config.sp_ulysses_degree == 4 + assert denoising.parallel_config.enable_fsdp + assert denoising.offload_config.offload_type == WeightOffloadType.NO_CPU_OFFLOAD + assert config.text_encoding_config.offload_config.offload_type == WeightOffloadType.MODEL_CPU_OFFLOAD + + +def test_build_config_rejects_sparse_attention_and_invalid_sp_degree() -> None: + for parallelism in (0, 3, 33): + try: + build_ltx25_distilled_config("cuda", torch.bfloat16, "diff", "none", parallelism=parallelism) + except ValueError: + pass + else: + raise AssertionError(f"parallelism={parallelism} should be rejected") + + try: + build_ltx25_distilled_config( + "cuda", + torch.bfloat16, + "diff", + "none", + attn_impl=AttnImplType.RADIAL_ATTN, + ) + except ValueError: + pass + else: + raise AssertionError("sparse attention should be rejected") + + +def test_i2v_run_with_file_accepts_service_image_alias(monkeypatch: pytest.MonkeyPatch) -> None: + opened_paths = [] + image = SimpleNamespace(convert=lambda mode: mode) + + def stop_run(*_args: object, **_kwargs: object) -> None: + raise RuntimeError("stop") + + monkeypatch.setattr(i2v_example.Image, "open", lambda path: opened_paths.append(path) or image) + monkeypatch.setattr(i2v_example, "run", stop_run) + + with pytest.raises(RuntimeError, match="stop"): + i2v_example.run_with_file( + object(), + prompt="test", + output_path="result.mp4", + first_image_path="service-input.png", + ) + + assert opened_paths == ["service-input.png"] diff --git a/tests/unit/pipelines/ltx25_distilled/test_reference.py b/tests/unit/pipelines/ltx25_distilled/test_reference.py new file mode 100644 index 0000000..6194646 --- /dev/null +++ b/tests/unit/pipelines/ltx25_distilled/test_reference.py @@ -0,0 +1,230 @@ +"""Monolithic LTX-2.5 distilled reference-path contracts.""" + +from __future__ import annotations + +import torch +from PIL import Image + +from telefuser.models.ltx25.embeddings import LTX25EmbeddingsProcessorOutput +from telefuser.models.ltx25.sampler import ancestral_noise_generator +from telefuser.pipelines.ltx25_distilled.reference import ( + LTX25DistilledReference, + LTX25ReferenceComponents, + LTX25ReferenceImageCondition, + LTX25ReferenceRequest, + _LazyCallable, + _LazyTextEncoder, + _release_modules, +) + + +def test_lazy_callable_proxies_instrumentation_attributes() -> None: + class Model: + velocity_model = object() + + lazy = _LazyCallable(lambda: Model()) + + assert lazy.velocity_model is not None + + +def test_release_modules_unloads_retained_lazy_callable() -> None: + class Model(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.to_calls = 0 + + def to(self, *args: object, **kwargs: object) -> "Model": + self.to_calls += 1 + return self + + model = Model() + lazy = _LazyCallable(lambda: model) + assert lazy.resolve() is model + + _release_modules(lazy) + + assert model.to_calls == 1 + assert lazy._model is None + + +def test_lazy_text_encoder_can_remain_resident_without_cpu_offload() -> None: + class Model(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.to_calls = 0 + + def encode(self, prompts: list[str]) -> tuple[tuple[torch.Tensor, ...], torch.Tensor, torch.Tensor]: + assert prompts == ["test"] + return (), torch.empty((1, 0), dtype=torch.long), torch.empty((1, 0), dtype=torch.long) + + def to(self, *args: object, **kwargs: object) -> "Model": + self.to_calls += 1 + return self + + model = Model() + loads = 0 + + def load() -> Model: + nonlocal loads + loads += 1 + return model + + encoder = _LazyTextEncoder(load, release_after_call=False) + + encoder.encode(["test"]) + encoder.encode(["test"]) + + assert loads == 1 + assert model.to_calls == 0 + + +class _TextEncoder: + def encode(self, prompts: list[str]) -> tuple[tuple[torch.Tensor, ...], torch.Tensor, torch.Tensor]: + assert prompts == ["A test prompt"] + return (torch.ones((1, 2, 4)),), torch.tensor([[1, 2]]), torch.tensor([[1, 1]]) + + +class _EmbeddingsProcessor: + def __call__( + self, hidden_states: tuple[torch.Tensor, ...], attention_mask: torch.Tensor + ) -> LTX25EmbeddingsProcessorOutput: + assert len(hidden_states) == 1 + return LTX25EmbeddingsProcessorOutput( + video_encoding=torch.ones((1, 2, 4)), + audio_encoding=torch.ones((1, 2, 4)), + attention_mask=attention_mask, + ) + + +class _Transformer: + def __call__(self, video: object, audio: object, perturbations: object) -> tuple[torch.Tensor, torch.Tensor]: + del perturbations + return torch.zeros_like(video.latent), torch.zeros_like(audio.latent) # type: ignore[union-attr] + + +class _RecordingTransformer(_Transformer): + def __init__(self) -> None: + self.video_latents: list[torch.Tensor] = [] + + def __call__(self, video: object, audio: object, perturbations: object) -> tuple[torch.Tensor, torch.Tensor]: + self.video_latents.append(video.latent.detach().clone()) # type: ignore[union-attr] + return super().__call__(video, audio, perturbations) + + +class _Upsampler: + def __call__(self, latent: torch.Tensor) -> torch.Tensor: + return latent.repeat_interleave(2, dim=3).repeat_interleave(2, dim=4) + + +class _IdentityStatistics: + def normalize(self, latent: torch.Tensor) -> torch.Tensor: + return latent + + def un_normalize(self, latent: torch.Tensor) -> torch.Tensor: + return latent + + +class _VideoEncoder: + def __init__(self) -> None: + self.pixel_shapes: list[torch.Size] = [] + + def __call__(self, pixels: torch.Tensor) -> torch.Tensor: + self.pixel_shapes.append(pixels.shape) + _, _, _, height, width = pixels.shape + return torch.ones((1, 128, 1, height // 32, width // 32), dtype=pixels.dtype) + + +def test_reference_preserves_two_stage_rng_and_latent_contracts() -> None: + reference = LTX25DistilledReference( + LTX25ReferenceComponents( + text_encoder=_TextEncoder(), + embeddings_processor=_EmbeddingsProcessor(), + transformer=_Transformer(), + spatial_upsampler=_Upsampler(), + latent_statistics=_IdentityStatistics(), + ), + device=torch.device("cpu"), + dtype=torch.bfloat16, + capture_prompt_intermediates=True, + ) + request = LTX25ReferenceRequest( + prompt="A test prompt", + seed=42, + height=256, + width=384, + num_frames=9, + frame_rate=24.0, + ) + + result = reference.generate(request) + + assert result.stage1_video.latent.shape == (1, 128, 2, 4, 6) + assert result.stage1_audio.latent.shape == (1, 8, 9, 16) + assert result.stage2_video.latent.shape == (1, 128, 2, 8, 12) + assert result.stage2_audio.latent.shape == (1, 8, 9, 16) + assert result.trace.video_context.shape == (1, 2, 4) + assert torch.equal(result.trace.artifacts["gemma_token_ids"], torch.tensor([[1, 2]])) + assert torch.equal(result.trace.artifacts["gemma_attention_mask"], torch.tensor([[1, 1]])) + assert "gemma_hidden_state_0" in result.trace.artifacts + for step_index in range(3): + assert f"stage2_step{step_index}_updated_video_latent" in result.trace.artifacts + assert f"stage2_step{step_index}_updated_audio_latent" in result.trace.artifacts + + initial_generator = torch.Generator().manual_seed(request.seed) + expected_stage1_video = torch.randn((1, 48, 128), generator=initial_generator, dtype=torch.bfloat16) + expected_stage1_audio = torch.randn((1, 9, 128), generator=initial_generator, dtype=torch.bfloat16) + torch.testing.assert_close(result.trace.artifacts["stage1_initial_video_noise"], expected_stage1_video) + torch.testing.assert_close(result.trace.artifacts["stage1_initial_audio_noise"], expected_stage1_audio) + + expected_stage2_video_noise = torch.randn((1, 192, 128), generator=initial_generator, dtype=torch.bfloat16) + expected_stage2_audio_noise = torch.randn((1, 9, 128), generator=initial_generator, dtype=torch.bfloat16) + expected_stage2_video = expected_stage2_video_noise * torch.tensor(0.909375, dtype=torch.bfloat16) + expected_stage2_audio = expected_stage2_audio_noise * torch.tensor(0.909375, dtype=torch.bfloat16) + torch.testing.assert_close(result.trace.artifacts["stage2_initial_video_noise"], expected_stage2_video) + torch.testing.assert_close(result.trace.artifacts["stage2_initial_audio_noise"], expected_stage2_audio) + assert torch.equal(result.decoder_generator_state, initial_generator.get_state()) + + ancestral = ancestral_noise_generator(request.seed, "cpu") + expected_ancestral_video = torch.randn((1, 48, 128), generator=ancestral, dtype=torch.bfloat16) + expected_ancestral_audio = torch.randn((1, 9, 128), generator=ancestral, dtype=torch.bfloat16) + torch.testing.assert_close(result.trace.artifacts["stage1_step0_ancestral_noise_video"], expected_ancestral_video) + torch.testing.assert_close(result.trace.artifacts["stage1_step0_ancestral_noise_audio"], expected_ancestral_audio) + + +def test_reference_applies_i2v_conditions_before_each_stage_noising() -> None: + encoder = _VideoEncoder() + transformer = _RecordingTransformer() + reference = LTX25DistilledReference( + LTX25ReferenceComponents( + text_encoder=_TextEncoder(), + embeddings_processor=_EmbeddingsProcessor(), + transformer=transformer, + spatial_upsampler=_Upsampler(), + latent_statistics=_IdentityStatistics(), + video_encoder=encoder, + ), + device=torch.device("cpu"), + dtype=torch.bfloat16, + ) + request = LTX25ReferenceRequest( + prompt="A test prompt", + seed=42, + height=256, + width=384, + num_frames=9, + frame_rate=24.0, + images=(LTX25ReferenceImageCondition(Image.new("RGB", (4, 4), "white"), crf=0),), + ) + + result = reference.generate(request) + + assert encoder.pixel_shapes == [torch.Size((1, 3, 1, 128, 192)), torch.Size((1, 3, 1, 256, 384))] + torch.testing.assert_close( + transformer.video_latents[1][:, :24], torch.ones_like(transformer.video_latents[1][:, :24]) + ) + torch.testing.assert_close( + result.stage1_video.latent[:, :, 0], torch.ones_like(result.stage1_video.latent[:, :, 0]) + ) + torch.testing.assert_close( + result.stage2_video.latent[:, :, 0], torch.ones_like(result.stage2_video.latent[:, :, 0]) + ) diff --git a/tests/unit/pipelines/ltx25_distilled/test_sampler.py b/tests/unit/pipelines/ltx25_distilled/test_sampler.py new file mode 100644 index 0000000..7db925f --- /dev/null +++ b/tests/unit/pipelines/ltx25_distilled/test_sampler.py @@ -0,0 +1,49 @@ +"""LTX-2.5 distilled sampler contract tests.""" + +from __future__ import annotations + +import pytest +import torch + +from telefuser.models.ltx25.sampler import ( + LTX25_STAGE1_DISTILLED_SIGMAS, + LTX25_STAGE2_DISTILLED_SIGMAS, + LTX25EulerAncestralStep, + ancestral_noise_generator, + distilled_sigmas, + uses_ancestral_stage1_sampler, +) + + +def test_distilled_schedules_are_frozen_upstream_values() -> None: + assert tuple(distilled_sigmas(1).tolist()) == pytest.approx(LTX25_STAGE1_DISTILLED_SIGMAS) + assert tuple(distilled_sigmas(2).tolist()) == pytest.approx(LTX25_STAGE2_DISTILLED_SIGMAS) + + +def test_ancestral_selection_starts_at_ltx_25() -> None: + assert not uses_ancestral_stage1_sampler("2.4") + assert uses_ancestral_stage1_sampler("2.5") + assert uses_ancestral_stage1_sampler("2.5-rc1") + + +def test_ancestral_step_matches_rectified_flow_equation() -> None: + sample = torch.tensor([2.0, -1.0]) + denoised = torch.tensor([1.0, 3.0]) + noise = torch.tensor([0.5, -0.25]) + sigmas = torch.tensor([1.0, 0.5]) + actual = LTX25EulerAncestralStep().step(sample, denoised, sigmas, 0, noise) + + expected = 0.25 * sample + 0.75 * denoised + expected = (0.5 / 0.75) * expected + noise * torch.sqrt(torch.tensor(0.5**2 - 0.25**2 * 0.5**2 / 0.75**2)) + torch.testing.assert_close(actual, expected) + + +def test_terminal_step_returns_denoised_without_noise() -> None: + result = LTX25EulerAncestralStep().step(torch.ones(2), torch.tensor([3.0, 4.0]), torch.tensor([0.5, 0.0]), 0) + torch.testing.assert_close(result, torch.tensor([3.0, 4.0])) + + +def test_ancestral_generator_is_offset_from_initial_seed() -> None: + actual = torch.randn(4, generator=ancestral_noise_generator(7, "cpu")) + expected = torch.randn(4, generator=torch.Generator().manual_seed(10_007)) + torch.testing.assert_close(actual, expected) diff --git a/tests/unit/pipelines/ltx25_distilled/test_video_decoding.py b/tests/unit/pipelines/ltx25_distilled/test_video_decoding.py new file mode 100644 index 0000000..befad64 --- /dev/null +++ b/tests/unit/pipelines/ltx25_distilled/test_video_decoding.py @@ -0,0 +1,30 @@ +"""LTX-2.5 video-decoding stage contracts.""" + +from __future__ import annotations + +import torch + +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.pipelines.ltx25_distilled.video_decoding import LTX25VideoDecodingStage + + +class _ConvVideoDecoder(torch.nn.Module): + def decode(self, latent: torch.Tensor, generator: torch.Generator) -> tuple[torch.Tensor, ...]: + del generator + return (latent[:, :3],) + + +def test_conv_vae_stage_converts_video_chunks_to_public_rgb_layout() -> None: + manager = ModuleManager(device="cpu", torch_dtype=torch.float32) + manager.add_module(_ConvVideoDecoder(), "ltx25_video_decoder") + stage = LTX25VideoDecodingStage( + manager, + ModelRuntimeConfig(device_type="cpu", torch_dtype=torch.float32), + video_vae="conv", + ) + + chunks = stage.decode(torch.zeros(1, 128, 2, 8, 12), torch.Generator(device="cpu")) + + assert chunks[0].shape == (2, 8, 12, 3) + assert torch.all(chunks[0] == 0.5) diff --git a/tests/unit/pipelines/minimax_h3/test_examples.py b/tests/unit/pipelines/minimax_h3/test_examples.py index 4655763..05ea7ab 100644 --- a/tests/unit/pipelines/minimax_h3/test_examples.py +++ b/tests/unit/pipelines/minimax_h3/test_examples.py @@ -5,6 +5,7 @@ from examples.minimax_h3 import minimax_h3_fl2va_h100 as fl2va_example from examples.minimax_h3 import minimax_h3_ref2va_h100 as ref2va_example +from examples.minimax_h3 import minimax_h3_turbo_lora_h100 as turbo_example from examples.minimax_h3.common import ( MINIMAX_H3_DEFAULT_FL2VA_IMAGE, MINIMAX_H3_DEFAULT_REF2VA_AUDIO, @@ -246,6 +247,33 @@ def __call__(self, **kwargs: object) -> object: assert [item["frame_index"] for item in calls[-1]["conditions"]] == [0, -1] +def test_turbo_run_uses_input_image_path(monkeypatch: pytest.MonkeyPatch) -> None: + calls = [] + marker = object() + + class Pipeline: + def __call__(self, **kwargs: object) -> object: + calls.append(kwargs) + return marker + + monkeypatch.setattr(turbo_example, "save_generation", lambda *_: None) + assert turbo_example.run(Pipeline(), input_image_path="input.png") is marker + assert calls[0]["conditions"] == [{"type": "image", "role": "keyframe", "uri": "input.png", "frame_index": 0}] + + +def test_turbo_run_with_file_accepts_service_image_alias(monkeypatch: pytest.MonkeyPatch) -> None: + calls = [] + + def fake_run(*args: object, **kwargs: object) -> object: + calls.append((args, kwargs)) + return object() + + monkeypatch.setattr(turbo_example, "run", fake_run) + turbo_example.run_with_file(object(), first_image_path="service-input.png", output_path="result.mp4") + + assert calls[0][1]["input_image_path"] == "service-input.png" + + def test_ref2va_run_preserves_ordered_service_conditions() -> None: calls = [] marker = object() diff --git a/tests/unit/service/test_example_service_parity.py b/tests/unit/service/test_example_service_parity.py index 8de6f59..dde078f 100644 --- a/tests/unit/service/test_example_service_parity.py +++ b/tests/unit/service/test_example_service_parity.py @@ -30,6 +30,9 @@ "wan21_i2v_service": (Path("examples/wan_video/wan21_14b_image_to_video_480p_service.py"), "i2v", True), "minimax_h3_fl2va": (Path("examples/minimax_h3/minimax_h3_fl2va_h100.py"), "t2v", True), "minimax_h3_ref2va": (Path("examples/minimax_h3/minimax_h3_ref2va_h100.py"), "s2v", True), + "minimax_h3_turbo": (Path("examples/minimax_h3/minimax_h3_turbo_lora_h100.py"), "i2v", True), + "ltx25_t2v": (Path("examples/ltx25_distilled/ltx25_distilled_t2v_h100.py"), "t2v", False), + "ltx25_i2v": (Path("examples/ltx25_distilled/ltx25_distilled_i2v_h100.py"), "i2v", False), "wan22_i2v_distill": (Path("examples/wan_video/wan22_14b_image_to_video_distill_h100.py"), "i2v", True), "lingbot_video_dense": (Path("examples/lingbot_video/lingbot_video_dense_1_3b.py"), "t2i", True), "lingbot_video_moe": (Path("examples/lingbot_video/lingbot_video_moe_30b.py"), "t2i", True), diff --git a/tests/unit/test_run_examples.py b/tests/unit/test_run_examples.py index 17f5a2e..72e5b24 100644 --- a/tests/unit/test_run_examples.py +++ b/tests/unit/test_run_examples.py @@ -197,6 +197,55 @@ def run_with_file(pipeline: object, output_path: str, target_video_length: float ) == {"output_path": "/tmp/result.mp4"} +def test_call_run_injects_configured_i2v_input_image_path(tmp_path: Path) -> None: + image_path = tmp_path / "input.png" + image_path.write_bytes(b"image") + module = ModuleType("i2v_example") + + def run_with_file( + pipeline: object, + input_image_path: str, + prompt: str, + output_path: str, + num_frames: int, + ) -> tuple[object, str, str, str, int]: + return pipeline, input_image_path, prompt, output_path, num_frames + + module.run_with_file = run_with_file + result = run_examples._call_run( + module, + "pipeline", + { + "input_image_path": str(image_path), + "prompt": "test prompt", + "output_path": "/tmp/result.mp4", + "num_frames": 121, + }, + entrypoint="run_with_file", + ) + + assert result == ("pipeline", str(image_path), "test prompt", "/tmp/result.mp4", 121) + + +def test_ltx25_two_gpu_regressions_are_registered() -> None: + config = run_examples.load_config() + + t2v = config.pipelines["ltx25_distilled_t2v_2gpu"] + assert t2v.script == "ltx25_distilled/ltx25_distilled_t2v_h100.py" + assert t2v.gpu_count == 2 + assert (t2v.width, t2v.height) == (1536, 1024) + assert t2v.use_run_with_file and t2v.require_audio + assert t2v.ppl_config_overrides["num_frames"] == 121 + + i2v = config.pipelines["ltx25_distilled_i2v_2gpu"] + assert i2v.script == "ltx25_distilled/ltx25_distilled_i2v_h100.py" + assert i2v.gpu_count == 2 + assert (i2v.width, i2v.height) == (896, 512) + assert i2v.input_image_path == "examples/data/ltx25/official_guitar_man.png" + assert i2v.use_run_with_file and i2v.require_audio + assert i2v.ppl_config_overrides["num_frames"] == 121 + + def test_video_metrics_rejects_mismatched_frame_counts(monkeypatch: pytest.MonkeyPatch) -> None: from telefuser.utils import video as video_utils diff --git a/tests/unit/validation/test_benchmark_ltx25_telefuser.py b/tests/unit/validation/test_benchmark_ltx25_telefuser.py new file mode 100644 index 0000000..07767bf --- /dev/null +++ b/tests/unit/validation/test_benchmark_ltx25_telefuser.py @@ -0,0 +1,18 @@ +"""Tests for LTX-2.5 benchmark report aggregation.""" + +from __future__ import annotations + +from tools.validation.benchmark_ltx25_telefuser import summarize_samples + + +def test_summarize_samples_preserves_raw_samples_and_reports_p50() -> None: + samples = [{"seconds": 3.0}, {"seconds": 1.0}, {"seconds": 2.0}] + + result = summarize_samples(samples) + + assert result["samples"] == samples + assert result["count"] == 3 + assert result["min_seconds"] == 1.0 + assert result["max_seconds"] == 3.0 + assert result["mean_seconds"] == 2.0 + assert result["p50_seconds"] == 2.0 diff --git a/tests/unit/validation/test_capture_ltx25_telefuser.py b/tests/unit/validation/test_capture_ltx25_telefuser.py new file mode 100644 index 0000000..f3bb615 --- /dev/null +++ b/tests/unit/validation/test_capture_ltx25_telefuser.py @@ -0,0 +1,14 @@ +"""Tests for TeleFuser LTX-2.5 capture helpers.""" + +from __future__ import annotations + +from tools.validation.capture_ltx25_telefuser import _release_modules + + +def test_capture_release_modules_releases_lazy_proxy() -> None: + class LazyProxy: + def __init__(self) -> None: + self.calls = 0 + + def release(self) -> None: + self.calls += 1 diff --git a/tests/unit/validation/test_compare_ltx25_benchmarks.py b/tests/unit/validation/test_compare_ltx25_benchmarks.py new file mode 100644 index 0000000..cbc23a7 --- /dev/null +++ b/tests/unit/validation/test_compare_ltx25_benchmarks.py @@ -0,0 +1,54 @@ +"""Tests for the LTX-2.5 benchmark comparison gate.""" + +from __future__ import annotations + +from tools.validation.compare_ltx25_benchmarks import compare_benchmarks + + +def _report(implementation: str, *, cold: float, warm: float, count: int = 5) -> dict[str, object]: + return { + "implementation": implementation, + "request": {"seed": 42, "offload": "cpu"}, + "runtime": { + "torch_version": "2.11.0", + "cuda_version": "12.8", + "natten_version": "0.21.6", + "natten_has_libnatten": True, + "gpus": [{"name": "H100"}], + }, + "cold": {"end_to_end": {"p50_seconds": cold, "count": count}}, + "warm": {"p50_seconds": warm, "count": count}, + } + + +def test_compare_benchmarks_passes_only_for_clear_speedup() -> None: + upstream = _report("upstream", cold=10.0, warm=8.0) + candidate = _report("telefuser", cold=9.7, warm=7.7) + + result = compare_benchmarks(upstream, candidate) + + assert result["passed"] + assert result["measurements"]["cold_end_to_end"]["status"] == "passed" + assert result["measurements"]["warm_end_to_end"]["status"] == "passed" + + +def test_compare_benchmarks_marks_noise_band_inconclusive() -> None: + upstream = _report("upstream", cold=10.0, warm=8.0) + candidate = _report("telefuser", cold=10.1, warm=8.0) + + result = compare_benchmarks(upstream, candidate) + + assert not result["passed"] + assert result["measurements"]["cold_end_to_end"]["status"] == "inconclusive" + + +def test_compare_benchmarks_rejects_mismatched_runtime_or_sample_count() -> None: + upstream = _report("upstream", cold=10.0, warm=8.0) + candidate = _report("telefuser", cold=9.0, warm=7.0, count=4) + candidate["runtime"] = {**candidate["runtime"], "natten_has_libnatten": False} # type: ignore[index] + + result = compare_benchmarks(upstream, candidate) + + assert not result["passed"] + assert not result["sufficient_samples"] + assert not result["runtime_match"]["natten_has_libnatten"] diff --git a/tests/unit/validation/test_ltx25_capture_utils.py b/tests/unit/validation/test_ltx25_capture_utils.py new file mode 100644 index 0000000..f513690 --- /dev/null +++ b/tests/unit/validation/test_ltx25_capture_utils.py @@ -0,0 +1,65 @@ +"""Tests for LTX-2.5 validation capture runtime controls.""" + +from __future__ import annotations + +from pathlib import Path + +import av +import numpy as np +import torch + +from tools.validation.ltx25_capture_utils import deterministic_audio_kernels, mp4_container_metadata + + +def test_deterministic_audio_kernels_restores_runtime_state() -> None: + """The capture-only mode must not leak process-wide CUDA settings.""" + benchmark = torch.backends.cudnn.benchmark + deterministic = torch.backends.cudnn.deterministic + deterministic_algorithms = torch.are_deterministic_algorithms_enabled() + try: + with deterministic_audio_kernels(True): + assert not torch.backends.cudnn.benchmark + assert torch.backends.cudnn.deterministic + assert torch.are_deterministic_algorithms_enabled() + finally: + torch.backends.cudnn.benchmark = benchmark + torch.backends.cudnn.deterministic = deterministic + torch.use_deterministic_algorithms(deterministic_algorithms) + assert torch.backends.cudnn.benchmark is benchmark + assert torch.backends.cudnn.deterministic is deterministic + assert torch.are_deterministic_algorithms_enabled() is deterministic_algorithms + + +def test_mp4_container_metadata_records_video_stream(tmp_path: Path) -> None: + """Container provenance includes the fields needed to replay a Golden output.""" + path = tmp_path / "decoded.mp4" + with av.open(str(path), mode="w") as container: + stream = container.add_stream("mpeg4", rate=24) + stream.width = 2 + stream.height = 2 + stream.pix_fmt = "yuv420p" + audio_stream = container.add_stream("aac", rate=48000) + audio_stream.layout = "stereo" + frame = av.VideoFrame.from_ndarray(np.zeros((2, 2, 3), dtype=np.uint8), format="rgb24") + audio_frame = av.AudioFrame.from_ndarray(np.zeros((2, 1024), dtype=np.float32), format="fltp", layout="stereo") + audio_frame.sample_rate = 48000 + for packet in stream.encode(frame): + container.mux(packet) + for packet in audio_stream.encode(audio_frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + for packet in audio_stream.encode(): + container.mux(packet) + + metadata = mp4_container_metadata(path) + + assert metadata["path"] == "decoded.mp4" + assert metadata["size_bytes"] > 0 + assert len(metadata["sha256"]) == 64 + video_stream = next(stream for stream in metadata["streams"] if stream["type"] == "video") + audio_stream = next(stream for stream in metadata["streams"] if stream["type"] == "audio") + assert video_stream["width"] == 2 + assert video_stream["height"] == 2 + assert audio_stream["sample_rate"] == 48000 + assert audio_stream["channels"] == 2 diff --git a/tools/validation/benchmark_ltx25_telefuser.py b/tools/validation/benchmark_ltx25_telefuser.py new file mode 100644 index 0000000..b285b4f --- /dev/null +++ b/tools/validation/benchmark_ltx25_telefuser.py @@ -0,0 +1,165 @@ +"""Record synchronized cold and warm LTX-2.5 TeleFuser pipeline measurements.""" + +from __future__ import annotations + +import argparse +import json +import statistics +from pathlib import Path +from typing import Any, Callable + +import torch +from PIL import Image + +from telefuser.metrics.runtime import collect_runtime_environment, finish_runtime_measurement, start_runtime_measurement +from telefuser.pipelines.ltx25_distilled import LTX25DistilledPipeline, LTX25ImageCondition + + +def _measure(device: torch.device, operation: Callable[[], Any]) -> tuple[Any, dict[str, Any]]: + """Execute one operation with synchronized timing and allocator peaks.""" + measurement = start_runtime_measurement([device], capture_peak_memory=True) + with torch.inference_mode(): + result = operation() + return result, finish_runtime_measurement(measurement) + + +def summarize_samples(samples: list[dict[str, Any]]) -> dict[str, Any]: + """Return stable raw timings plus p50 summary for one measured phase.""" + seconds = [float(sample["seconds"]) for sample in samples] + if not seconds: + raise ValueError("benchmark requires at least one measured sample") + return { + "samples": samples, + "count": len(seconds), + "min_seconds": min(seconds), + "max_seconds": max(seconds), + "mean_seconds": statistics.mean(seconds), + "p50_seconds": statistics.median(seconds), + } + + +def _request_image(path: Path | None, frame_index: int, strength: float) -> tuple[LTX25ImageCondition, ...]: + if path is None: + return () + if not path.is_file(): + raise FileNotFoundError(f"LTX-2.5 conditioning image does not exist: {path}") + return (LTX25ImageCondition(Image.open(path).convert("RGB"), frame_idx=frame_index, strength=strength),) + + +def _request_kwargs(args: argparse.Namespace, images: tuple[LTX25ImageCondition, ...]) -> dict[str, Any]: + return { + "prompt": args.prompt, + "seed": args.seed, + "height": args.height, + "width": args.width, + "num_frames": args.num_frames, + "frame_rate": args.frame_rate, + "images": images, + } + + +def _build_pipeline(args: argparse.Namespace) -> LTX25DistilledPipeline: + return LTX25DistilledPipeline.from_model_root( + args.model_root, + device=args.device, + video_vae=args.video_vae, + offload=args.offload, + ) + + +def main() -> None: + """Run the requested cold/warm benchmark and write one provenance-rich report.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--prompt", required=True) + parser.add_argument("--seed", type=int, required=True) + parser.add_argument("--height", type=int, default=1024) + parser.add_argument("--width", type=int, default=1536) + parser.add_argument("--num-frames", type=int, default=121) + parser.add_argument("--frame-rate", type=float, default=24.0) + parser.add_argument("--video-vae", choices=("diff", "conv"), default="diff") + parser.add_argument("--offload", choices=("none", "cpu"), required=True) + parser.add_argument("--image", type=Path) + parser.add_argument("--image-frame-index", type=int, default=0) + parser.add_argument("--image-strength", type=float, default=1.0) + parser.add_argument("--device", default="cuda") + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--runs", type=int, default=5) + args = parser.parse_args() + if args.warmup < 0 or args.runs < 1: + raise ValueError("--warmup must be non-negative and --runs must be positive") + + device = torch.device(args.device) + if device.type != "cuda" or not torch.cuda.is_available(): + raise RuntimeError("LTX-2.5 formal benchmarking requires a CUDA device") + images = _request_image(args.image, args.image_frame_index, args.image_strength) + request = _request_kwargs(args, images) + + cold_construction: list[dict[str, Any]] = [] + cold_generation: list[dict[str, Any]] = [] + for _ in range(args.runs): + cold_pipeline, construction = _measure(device, lambda: _build_pipeline(args)) + _, generation = _measure(device, lambda pipeline=cold_pipeline: pipeline(**request)) + cold_construction.append(construction) + cold_generation.append(generation) + del cold_pipeline + torch.cuda.empty_cache() + + warm_result = _measure(device, lambda: _build_pipeline(args)) + warm_pipeline, warm_construction = warm_result + warmup: list[dict[str, Any]] = [] + for _ in range(args.warmup): + _, measurement = _measure(device, lambda pipeline=warm_pipeline: pipeline(**request)) + warmup.append(measurement) + samples: list[dict[str, Any]] = [] + for _ in range(args.runs): + _, measurement = _measure(device, lambda pipeline=warm_pipeline: pipeline(**request)) + samples.append(measurement) + + try: + import natten + + natten_capability: bool | None = bool(getattr(natten, "HAS_LIBNATTEN", False)) + natten_version: str | None = getattr(natten, "__version__", None) + except ImportError: + natten_capability, natten_version = None, None + report = { + "implementation": "telefuser", + "runtime": { + **collect_runtime_environment([device], repo_root=Path.cwd()), + "natten_version": natten_version, + "natten_has_libnatten": natten_capability, + }, + "request": { + "prompt": args.prompt, + "seed": args.seed, + "height": args.height, + "width": args.width, + "num_frames": args.num_frames, + "frame_rate": args.frame_rate, + "video_vae": args.video_vae, + "offload": args.offload, + "image": None + if args.image is None + else { + "path": str(args.image.resolve()), + "frame_index": args.image_frame_index, + "strength": args.image_strength, + }, + }, + "cold": { + "pipeline_construction": summarize_samples(cold_construction), + "end_to_end": summarize_samples(cold_generation), + }, + "warm_pipeline_construction": warm_construction, + "warmup": warmup, + "warm": summarize_samples(samples), + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/benchmark_ltx25_upstream.py b/tools/validation/benchmark_ltx25_upstream.py new file mode 100644 index 0000000..4781c55 --- /dev/null +++ b/tools/validation/benchmark_ltx25_upstream.py @@ -0,0 +1,181 @@ +"""Record synchronized cold and warm LTX-2.5 upstream pipeline measurements. + +Run with the pinned upstream LTX source on ``PYTHONPATH``. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +from pathlib import Path +from typing import Any, Callable + +import torch + +from telefuser.metrics.runtime import collect_runtime_environment, finish_runtime_measurement, start_runtime_measurement + + +def _measure(device: torch.device, operation: Callable[[], Any]) -> tuple[Any, dict[str, Any]]: + """Execute one operation with synchronized timing and allocator peaks.""" + measurement = start_runtime_measurement([device], capture_peak_memory=True) + with torch.inference_mode(): + result = operation() + return result, finish_runtime_measurement(measurement) + + +def _summarize_samples(samples: list[dict[str, Any]]) -> dict[str, Any]: + """Return stable raw timings plus p50 summary for one measured phase.""" + seconds = [float(sample["seconds"]) for sample in samples] + if not seconds: + raise ValueError("benchmark requires at least one measured sample") + return { + "samples": samples, + "count": len(seconds), + "min_seconds": min(seconds), + "max_seconds": max(seconds), + "mean_seconds": statistics.mean(seconds), + "p50_seconds": statistics.median(seconds), + } + + +def main() -> None: + """Run the requested upstream cold/warm benchmark and write one report.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--prompt", required=True) + parser.add_argument("--seed", type=int, required=True) + parser.add_argument("--height", type=int, default=1024) + parser.add_argument("--width", type=int, default=1536) + parser.add_argument("--num-frames", type=int, default=121) + parser.add_argument("--frame-rate", type=float, default=24.0) + parser.add_argument("--video-vae", choices=("diff", "conv"), default="diff") + parser.add_argument("--offload", choices=("none", "cpu"), required=True) + parser.add_argument("--image", type=Path) + parser.add_argument("--image-frame-index", type=int, default=0) + parser.add_argument("--image-strength", type=float, default=1.0) + parser.add_argument("--device", default="cuda") + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--runs", type=int, default=5) + args = parser.parse_args() + if args.warmup < 0 or args.runs < 1: + raise ValueError("--warmup must be non-negative and --runs must be positive") + if args.image is not None and not args.image.is_file(): + raise FileNotFoundError(f"LTX-2.5 conditioning image does not exist: {args.image}") + + from ltx_pipelines.distilled import DistilledPipeline # type: ignore[import-not-found] + from ltx_pipelines.utils.args import ImageConditioningInput # type: ignore[import-not-found] + from ltx_pipelines.utils.model_paths import ModelPaths # type: ignore[import-not-found] + from ltx_pipelines.utils.types import OffloadMode # type: ignore[import-not-found] + + device = torch.device(args.device) + if device.type != "cuda" or not torch.cuda.is_available(): + raise RuntimeError("LTX-2.5 formal benchmarking requires a CUDA device") + video_vae_name = ( + "ltx-2.5-video-vae-bf16.safetensors" if args.video_vae == "diff" else "ltx-2.5-video-vae-conv-bf16.safetensors" + ) + paths = ModelPaths.from_split( + transformer_path=str(args.model_root / "diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors"), + text_encoder_path=str(args.model_root / "text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors"), + video_vae_path=str(args.model_root / "vae" / video_vae_name), + audio_vae_path=str(args.model_root / "vae/ltx-2.5-audio-vae-bf16.safetensors"), + duration_head_path=str(args.model_root / "model_patches/ltx-2.5-duration-head-bf16.safetensors"), + ) + images = ( + [] + if args.image is None + else [ + ImageConditioningInput( + path=str(args.image.resolve()), frame_idx=args.image_frame_index, strength=args.image_strength + ) + ] + ) + request = { + "prompt": args.prompt, + "seed": args.seed, + "height": args.height, + "width": args.width, + "frame_rate": args.frame_rate, + "images": images, + "num_frames": args.num_frames, + } + + def build_pipeline() -> Any: + return DistilledPipeline( + model_paths=paths, + spatial_upsampler_path=str( + args.model_root / "latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors" + ), + loras=(), + device=device, + offload_mode=OffloadMode(args.offload), + ) + + cold_construction: list[dict[str, Any]] = [] + cold_generation: list[dict[str, Any]] = [] + for _ in range(args.runs): + cold_pipeline, construction = _measure(device, build_pipeline) + _, generation = _measure(device, lambda pipeline=cold_pipeline: pipeline(**request)) + cold_construction.append(construction) + cold_generation.append(generation) + del cold_pipeline + torch.cuda.empty_cache() + + warm_result = _measure(device, build_pipeline) + warm_pipeline, warm_construction = warm_result + warmup: list[dict[str, Any]] = [] + for _ in range(args.warmup): + _, measurement = _measure(device, lambda pipeline=warm_pipeline: pipeline(**request)) + warmup.append(measurement) + samples: list[dict[str, Any]] = [] + for _ in range(args.runs): + _, measurement = _measure(device, lambda pipeline=warm_pipeline: pipeline(**request)) + samples.append(measurement) + + try: + import natten + + natten_capability: bool | None = bool(getattr(natten, "HAS_LIBNATTEN", False)) + natten_version: str | None = getattr(natten, "__version__", None) + except ImportError: + natten_capability, natten_version = None, None + report = { + "implementation": "upstream", + "runtime": { + **collect_runtime_environment([device], repo_root=Path.cwd()), + "natten_version": natten_version, + "natten_has_libnatten": natten_capability, + }, + "request": { + "prompt": args.prompt, + "seed": args.seed, + "height": args.height, + "width": args.width, + "num_frames": args.num_frames, + "frame_rate": args.frame_rate, + "video_vae": args.video_vae, + "offload": args.offload, + "image": None + if args.image is None + else { + "path": str(args.image.resolve()), + "frame_index": args.image_frame_index, + "strength": args.image_strength, + }, + }, + "cold": { + "pipeline_construction": _summarize_samples(cold_construction), + "end_to_end": _summarize_samples(cold_generation), + }, + "warm_pipeline_construction": warm_construction, + "warmup": warmup, + "warm": _summarize_samples(samples), + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/capture_ltx25_telefuser.py b/tools/validation/capture_ltx25_telefuser.py new file mode 100644 index 0000000..a4bfc47 --- /dev/null +++ b/tools/validation/capture_ltx25_telefuser.py @@ -0,0 +1,349 @@ +"""Capture TeleFuser LTX-2.5 distilled T2V stage artifacts for golden comparison.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path +from typing import Any + +import torch +from PIL import Image + +try: + from ltx25_capture_utils import deterministic_audio_kernels +except ModuleNotFoundError: # Supports runpy-based capture wrappers from the repository root. + from tools.validation.ltx25_capture_utils import deterministic_audio_kernels + +from telefuser.models.ltx25.checkpoint import LTX25ModelPaths, inspect_model_pack +from telefuser.models.ltx25.sampler import distilled_sigmas +from telefuser.pipelines.ltx25_distilled.reference import ( + LTX25DistilledReference, + LTX25ReferenceImageCondition, + LTX25ReferenceRequest, +) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _tiling_config_metadata(tiling_config: Any) -> dict[str, dict[str, int]]: + """Serialize resolved DiffVAE tile geometry for a reproducible capture manifest.""" + return { + axis: {"tile_size": config.tile_size, "overlap": config.overlap} + for axis, config in ( + ("frames", tiling_config.frames), + ("height", tiling_config.height), + ("width", tiling_config.width), + ) + } + + +def _save_tensor(output_dir: Path, name: str, tensor: torch.Tensor) -> dict[str, Any]: + path = output_dir / f"{name}.pt" + value = tensor.detach().cpu().contiguous() + torch.save(value, path) + return { + "path": path.name, + "sha256": _sha256(path), + "shape": list(value.shape), + "dtype": str(value.dtype), + } + + +def _diffvae_attention_backends(decoder: torch.nn.Module) -> list[str | None]: + """Return the NATTEN backend identities used by a loaded DiffVAE decoder.""" + backends = {module.natten_backend for module in decoder.modules() if hasattr(module, "natten_backend")} + return sorted(backends, key=lambda backend: backend or "") + + +def _capture_block_outputs( + transformer: torch.nn.Module, + output_dir: Path, + artifacts: dict[str, Any], + prefix: str, +) -> list[torch.utils.hooks.RemovableHandle]: + """Capture post-block video and audio states for one denoiser invocation.""" + velocity_model = getattr(transformer, "velocity_model") + handles: list[torch.utils.hooks.RemovableHandle] = [] + for index, block in enumerate(velocity_model.transformer_blocks): + + def save_output( + _module: torch.nn.Module, + _inputs: tuple[Any, ...], + output: tuple[Any, Any], + *, + block_index: int = index, + ) -> None: + video, audio = output + if video is not None: + name = f"{prefix}_block{block_index}_video" + artifacts[name] = _save_tensor(output_dir, name, video.x) + if audio is not None: + name = f"{prefix}_block{block_index}_audio" + artifacts[name] = _save_tensor(output_dir, name, audio.x) + + handles.append(block.register_forward_hook(save_output)) + return handles + + +def _release_modules(*modules: object, offload: str) -> None: + """Release reference-path modules before loading the decoders.""" + if offload == "cpu": + for module in modules: + release = getattr(module, "release", None) + if callable(release): + release() + elif isinstance(module, torch.nn.Module): + module.to("cpu") + if torch.cuda.is_available(): + torch.cuda.synchronize() + torch.cuda.empty_cache() + + +def _trajectories(artifacts: dict[str, Any], device: torch.device) -> dict[str, list[dict[str, Any]]]: + """Build the upstream-compatible per-step diffusion index from saved tensors.""" + trajectories: dict[str, list[dict[str, Any]]] = {} + for stage_name, stage_number in (("stage1", 1), ("stage2", 2)): + rows: list[dict[str, Any]] = [] + for step_index, sigma in enumerate(distilled_sigmas(stage_number, device=device)[:-1]): + row: dict[str, Any] = {"index": step_index, "sigma": float(sigma)} + for modality in ("video", "audio"): + updated = f"{stage_name}_step{step_index}_updated_{modality}_latent" + if updated in artifacts: + row[f"updated_{modality}_latent"] = updated + ancestral_noise = f"{stage_name}_step{step_index}_ancestral_noise_{modality}" + if ancestral_noise in artifacts: + row[f"ancestral_noise_{modality}"] = ancestral_noise + rows.append(row) + trajectories[stage_name] = rows + return trajectories + + +@torch.inference_mode() +def main() -> None: + """Run the isolated TeleFuser T2V reference path and save stage-boundary artifacts.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--prompt", required=True) + parser.add_argument("--seed", type=int, required=True) + parser.add_argument("--height", type=int, required=True) + parser.add_argument("--width", type=int, required=True) + parser.add_argument("--num-frames", type=int, required=True) + parser.add_argument("--frame-rate", type=float, default=24.0) + parser.add_argument("--video-vae", choices=("diff", "conv"), default="diff") + parser.add_argument("--image", type=Path) + parser.add_argument("--image-frame-index", type=int, default=0) + parser.add_argument("--image-strength", type=float, default=1.0) + parser.add_argument("--offload", choices=("none", "cpu"), default="cpu") + parser.add_argument("--device", default="cuda") + parser.add_argument("--capture-prompt-intermediates", action="store_true") + parser.add_argument("--capture-stage2-step2-blocks", action="store_true") + parser.add_argument("--capture-decoded", action="store_true") + parser.add_argument( + "--deterministic-audio", + action="store_true", + help="Use deterministic CUDA kernels only while capturing decoded audio.", + ) + args = parser.parse_args() + + args.output_dir.mkdir(parents=True, exist_ok=True) + if args.image is not None and not args.image.is_file(): + raise FileNotFoundError(f"LTX-2.5 conditioning image does not exist: {args.image}") + images = ( + () + if args.image is None + else ( + LTX25ReferenceImageCondition( + Image.open(args.image).convert("RGB"), + frame_idx=args.image_frame_index, + strength=args.image_strength, + ), + ) + ) + request = LTX25ReferenceRequest( + prompt=args.prompt, + seed=args.seed, + height=args.height, + width=args.width, + num_frames=args.num_frames, + frame_rate=args.frame_rate, + images=images, + ) + runner = LTX25DistilledReference.from_model_root( + str(args.model_root), + device=args.device, + video_vae=args.video_vae, + capture_prompt_intermediates=args.capture_prompt_intermediates, + offload=args.offload, + ) + block_artifacts: dict[str, Any] = {} + if args.capture_stage2_step2_blocks: + transformer = runner.components.transformer + original_forward = transformer.forward + calls = 0 + + def traced_forward(*forward_args: Any, **forward_kwargs: Any) -> Any: + nonlocal calls + handles: list[torch.utils.hooks.RemovableHandle] = [] + if calls == 10: + handles = _capture_block_outputs(transformer, args.output_dir, block_artifacts, "stage2_step2") + try: + return original_forward(*forward_args, **forward_kwargs) + finally: + for handle in handles: + handle.remove() + calls += 1 + + transformer.forward = traced_forward # type: ignore[method-assign] + try: + with torch.inference_mode(): + result = runner.generate(request) + finally: + transformer.forward = original_forward # type: ignore[method-assign] + else: + with torch.inference_mode(): + result = runner.generate(request) + + decoded_artifacts: dict[str, Any] = {} + resolved_video_tiling: dict[str, dict[str, int]] | None = None + resolved_diffvae_attention_backends: list[str | None] | None = None + resolved_diffvae_optimization: str | None = None + if args.capture_decoded: + from telefuser.models.ltx25 import ( + DiffusionVideoDecoder, + LTX25ConvVideoVAE, + load_ltx25_audio_decoder_and_vocoder, + ) + + _release_modules( + runner.components.text_encoder, + runner.components.embeddings_processor, + runner.components.transformer, + runner.components.spatial_upsampler, + runner.components.latent_statistics, + offload=args.offload, + ) + decoder_generator = torch.Generator(device=runner.device) + decoder_generator.set_state(result.decoder_generator_state) + paths = LTX25ModelPaths.from_model_root(args.model_root) + if args.video_vae == "diff": + video_decoder = DiffusionVideoDecoder.from_checkpoint( + paths.video_vae_path, + device=runner.device, + torch_dtype=runner.dtype, + ) + tiling_config = video_decoder.recommended_tiling_config( + height=args.height, + width=args.width, + num_frames=args.num_frames, + ) + resolved_video_tiling = _tiling_config_metadata(tiling_config) + resolved_diffvae_attention_backends = _diffvae_attention_backends(video_decoder) + resolved_diffvae_optimization = "chunked_compile" if video_decoder.mark_dynamic_shapes else "chunked_eager" + decoded_rgb = torch.cat( + list( + video_decoder.decode_video( + result.stage2_video.latent, tiling_config=tiling_config, generator=decoder_generator + ) + ), + dim=0, + ) + else: + video_decoder = LTX25ConvVideoVAE.from_checkpoint( + paths.conv_video_vae_path, + device=runner.device, + torch_dtype=runner.dtype, + ) + decoded_rgb = torch.cat( + [ + chunk[0].permute(1, 2, 3, 0).add(1).mul(0.5).clamp(0, 1) + for chunk in video_decoder.decode(result.stage2_video.latent, generator=decoder_generator) + ], + dim=0, + ) + decoded_artifacts["video_decoder_generator_state"] = _save_tensor( + args.output_dir, "video_decoder_generator_state", result.decoder_generator_state + ) + decoded_artifacts["video_decoder_input_latent"] = _save_tensor( + args.output_dir, "video_decoder_input_latent", result.stage2_video.latent + ) + decoded_artifacts["decoded_rgb"] = _save_tensor(args.output_dir, "decoded_rgb", decoded_rgb) + _release_modules(video_decoder, offload=args.offload) + + audio_decoder, vocoder = load_ltx25_audio_decoder_and_vocoder( + paths.audio_vae_path, + device=runner.device, + torch_dtype=runner.dtype, + ) + with deterministic_audio_kernels(args.deterministic_audio): + waveform = vocoder(audio_decoder(result.stage2_audio.latent)).squeeze(0).float() + decoded_artifacts["decoded_waveform"] = _save_tensor(args.output_dir, "decoded_waveform", waveform) + _release_modules(audio_decoder, vocoder, offload=args.offload) + + artifacts = { + name: _save_tensor(args.output_dir, name, value) for name, value in sorted(result.trace.artifacts.items()) + } + artifacts.update(block_artifacts) + artifacts.update(decoded_artifacts) + artifacts["video_context"] = _save_tensor(args.output_dir, "video_context", result.trace.video_context) + artifacts["audio_context"] = _save_tensor(args.output_dir, "audio_context", result.trace.audio_context) + artifacts["context_attention_mask"] = _save_tensor( + args.output_dir, "context_attention_mask", result.trace.context_attention_mask + ) + manifest = { + "runtime": { + "torch": torch.__version__, + "device": str(args.device), + "dtype": "bfloat16", + "deterministic_audio": args.deterministic_audio, + "tiling": "auto", + "resolved_video_tiling": resolved_video_tiling, + "resolved_diffvae_attention_backends": resolved_diffvae_attention_backends, + "resolved_diffvae_optimization": resolved_diffvae_optimization, + }, + "request": { + "prompt": request.prompt, + "seed": request.seed, + "height": request.height, + "width": request.width, + "num_frames": request.num_frames, + "resolved_num_frames": request.num_frames, + "frame_rate": request.frame_rate, + "video_vae": args.video_vae, + "offload": args.offload, + "dtype": "bfloat16", + "prompt_normalization": [request.prompt.strip()], + "image": ( + None + if args.image is None + else { + "path": str(args.image.resolve()), + "sha256": _sha256(args.image), + "frame_index": args.image_frame_index, + "strength": args.image_strength, + } + ), + }, + "checkpoints": { + name: metadata.as_dict() + for name, metadata in inspect_model_pack(args.model_root, include_sha256=True).items() + }, + "audio": {"sample_rate": 48000}, + "artifacts": artifacts, + "trajectories": _trajectories(artifacts, runner.device), + } + (args.output_dir / "capture_manifest.json").write_text( + json.dumps(manifest, indent=2, sort_keys=True), encoding="utf-8" + ) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/capture_ltx25_upstream.py b/tools/validation/capture_ltx25_upstream.py new file mode 100644 index 0000000..b1a2c8c --- /dev/null +++ b/tools/validation/capture_ltx25_upstream.py @@ -0,0 +1,517 @@ +"""Capture LTX-2.5 distilled upstream stage-boundary Golden artifacts. + +Run this tool with the pinned upstream source on ``PYTHONPATH``. The reference +runtime must use the same PyTorch and Transformers versions as TeleFuser. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import platform +import subprocess +import sys +from dataclasses import asdict +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path +from typing import Any + +import torch + +try: + from ltx25_capture_utils import deterministic_audio_kernels, mp4_container_metadata +except ModuleNotFoundError: # Supports runpy-based capture wrappers from the repository root. + from tools.validation.ltx25_capture_utils import deterministic_audio_kernels, mp4_container_metadata + +from telefuser.models.ltx25.checkpoint import inspect_model_pack + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _optional_package_version(package: str) -> str | None: + """Return an installed package version without making it a capture dependency.""" + try: + return version(package) + except PackageNotFoundError: + return None + + +def _save_tensor(output_dir: Path, name: str, tensor: torch.Tensor) -> dict[str, Any]: + path = output_dir / f"{name}.pt" + value = tensor.detach().cpu().contiguous() + torch.save(value, path) + return { + "path": path.name, + "sha256": _sha256(path), + "shape": list(value.shape), + "dtype": str(value.dtype), + } + + +def _capture_block_outputs( + transformer: torch.nn.Module, + output_dir: Path, + artifacts: dict[str, Any], + prefix: str, +) -> list[torch.utils.hooks.RemovableHandle]: + """Capture post-block video and audio states for one denoiser invocation.""" + velocity_model = getattr(transformer, "velocity_model") + handles: list[torch.utils.hooks.RemovableHandle] = [] + for index, block in enumerate(velocity_model.transformer_blocks): + + def save_output( + _module: torch.nn.Module, + _inputs: tuple[Any, ...], + output: tuple[Any, Any], + *, + block_index: int = index, + ) -> None: + video, audio = output + if video is not None: + name = f"{prefix}_block{block_index}_video" + artifacts[name] = _save_tensor(output_dir, name, video.x) + if audio is not None: + name = f"{prefix}_block{block_index}_audio" + artifacts[name] = _save_tensor(output_dir, name, audio.x) + + handles.append(block.register_forward_hook(save_output)) + return handles + + +class _StageRecorder: + """Record initial, per-step, and final upstream diffusion-stage tensors.""" + + def __init__( + self, + stage: Any, + output_dir: Path, + artifacts: dict[str, Any], + trajectories: dict[str, list[dict[str, Any]]], + capture_stage2_step2_blocks: bool, + ) -> None: + self._stage = stage + self._output_dir = output_dir + self._artifacts = artifacts + self._trajectories = trajectories + self._capture_stage2_step2_blocks = capture_stage2_step2_blocks + self._calls = 0 + + def __call__(self, *args: Any, **kwargs: Any) -> tuple[Any, Any]: + from ltx_core.components.diffusion_steps import EulerDiffusionStep # type: ignore[import-not-found] + + stage_name = f"stage{self._calls + 1}" + original_denoiser = kwargs["denoiser"] + stepper = kwargs.get("stepper") or EulerDiffusionStep() + kwargs["stepper"] = stepper + trajectories = self._trajectories.setdefault(stage_name, []) + + def record_denoiser( + transformer: Any, + video_state: Any, + audio_state: Any, + sigmas: torch.Tensor, + step_index: int, + ) -> tuple[Any, Any]: + if step_index == 0: + if video_state is not None: + self._artifacts[f"{stage_name}_initial_video_noise"] = _save_tensor( + self._output_dir, f"{stage_name}_initial_video_noise", video_state.latent + ) + if audio_state is not None: + self._artifacts[f"{stage_name}_initial_audio_noise"] = _save_tensor( + self._output_dir, f"{stage_name}_initial_audio_noise", audio_state.latent + ) + handles: list[torch.utils.hooks.RemovableHandle] = [] + if self._capture_stage2_step2_blocks and stage_name == "stage2" and step_index == 2: + handles = _capture_block_outputs(transformer, self._output_dir, self._artifacts, "stage2_step2") + try: + video_result, audio_result = original_denoiser( + transformer, video_state, audio_state, sigmas, step_index + ) + finally: + for handle in handles: + handle.remove() + row: dict[str, Any] = {"index": step_index, "sigma": float(sigmas[step_index])} + if video_result is not None and video_result.denoised is not None: + artifact = f"{stage_name}_step{step_index}_video_x0" + self._artifacts[artifact] = _save_tensor(self._output_dir, artifact, video_result.denoised) + if audio_result is not None and audio_result.denoised is not None: + artifact = f"{stage_name}_step{step_index}_audio_x0" + self._artifacts[artifact] = _save_tensor(self._output_dir, artifact, audio_result.denoised) + trajectories.append(row) + return video_result, audio_result + + original_step = stepper.step + + def record_step(*step_args: Any, **step_kwargs: Any) -> torch.Tensor: + step_index = step_kwargs["step_index"] if "step_index" in step_kwargs else step_args[3] + noise = step_kwargs.get("noise") + if noise is not None: + # The upstream loop draws video before audio from one generator. Recording here + # preserves both the actual tensor and that call order without changing RNG. + modality = "video" if "ancestral_noise_video" not in trajectories[step_index] else "audio" + artifact = f"{stage_name}_step{step_index}_ancestral_noise_{modality}" + self._artifacts[artifact] = _save_tensor(self._output_dir, artifact, noise) + trajectories[step_index][f"ancestral_noise_{modality}"] = artifact + updated = original_step(*step_args, **step_kwargs) + modality = "video" if "updated_video_latent" not in trajectories[step_index] else "audio" + artifact = f"{stage_name}_step{step_index}_updated_{modality}_latent" + self._artifacts[artifact] = _save_tensor(self._output_dir, artifact, updated) + trajectories[step_index][f"updated_{modality}_latent"] = artifact + return updated + + stepper.step = record_step + kwargs["denoiser"] = record_denoiser + video, audio = self._stage(*args, **kwargs) + self._artifacts[f"{stage_name}_video_latent"] = _save_tensor( + self._output_dir, f"{stage_name}_video_latent", video.latent + ) + self._artifacts[f"{stage_name}_audio_latent"] = _save_tensor( + self._output_dir, f"{stage_name}_audio_latent", audio.latent + ) + self._calls += 1 + return video, audio + + def __getattr__(self, name: str) -> Any: + return getattr(self._stage, name) + + +class _UpsamplerRecorder: + """Record the exact upsampler input and output used between denoising stages.""" + + def __init__(self, upsampler: Any, output_dir: Path, artifacts: dict[str, Any]) -> None: + self._upsampler = upsampler + self._output_dir = output_dir + self._artifacts = artifacts + + def __call__(self, latent: torch.Tensor) -> torch.Tensor: + self._artifacts["upsampler_input"] = _save_tensor(self._output_dir, "upsampler_input", latent) + output = self._upsampler(latent) + self._artifacts["upsampler_output"] = _save_tensor(self._output_dir, "upsampler_output", output) + return output + + def __getattr__(self, name: str) -> Any: + return getattr(self._upsampler, name) + + +class _VideoDecoderRecorder: + """Record the exact generator state entering the upstream video decoder.""" + + def __init__(self, decoder: Any, output_dir: Path, artifacts: dict[str, Any]) -> None: + self._decoder = decoder + self._output_dir = output_dir + self._artifacts = artifacts + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + latent = kwargs.get("latent") if "latent" in kwargs else args[0] + if isinstance(latent, torch.Tensor): + self._artifacts["video_decoder_input_latent"] = _save_tensor( + self._output_dir, "video_decoder_input_latent", latent + ) + generator = kwargs.get("generator") + if generator is None and len(args) >= 3: + generator = args[2] + if generator is not None: + self._artifacts["video_decoder_generator_state"] = _save_tensor( + self._output_dir, + "video_decoder_generator_state", + generator.get_state(), + ) + return self._decoder(*args, **kwargs) + + def __getattr__(self, name: str) -> Any: + return getattr(self._decoder, name) + + +class _PromptRecorder: + """Record raw Gemma inputs/states and final connectors without changing lifecycle.""" + + def __init__(self, encoder: Any, output_dir: Path, artifacts: dict[str, Any]) -> None: + self._encoder = encoder + self._output_dir = output_dir + self._artifacts = artifacts + self.normalized_prompts: list[str] = [] + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + from ltx_core.text_encoders.gemma.embeddings_processor import ( + EmbeddingsProcessor, # type: ignore[import-not-found] + ) + from ltx_core.text_encoders.gemma.encoders.base_encoder import ( + LTXGemmaTextEncoder, # type: ignore[import-not-found] + ) + + original_encode = LTXGemmaTextEncoder.encode + original_process_hidden_states = EmbeddingsProcessor.process_hidden_states + + def record_encode(text_encoder: Any, prompts: list[str]) -> Any: + if len(prompts) != 1: + raise ValueError(f"Golden capture expects one prompt, got {len(prompts)}") + self.normalized_prompts = [prompt.strip() for prompt in prompts] + pairs = text_encoder.tokenizer.tokenize_with_weights(prompts[0])["gemma"] + token_ids = torch.tensor([[token for token, _ in pairs]], device=text_encoder.model.device) + raw_mask = torch.tensor([[weight for _, weight in pairs]], device=text_encoder.model.device) + self._artifacts["gemma_token_ids"] = _save_tensor(self._output_dir, "gemma_token_ids", token_ids) + self._artifacts["gemma_attention_mask"] = _save_tensor(self._output_dir, "gemma_attention_mask", raw_mask) + raw_outputs = original_encode(text_encoder, prompts) + hidden_states, attention_mask = raw_outputs[0] + if not torch.equal(raw_mask, attention_mask): + raise ValueError("Golden capture token mask differs from the encoder's raw attention mask") + for index, hidden_state in enumerate(hidden_states): + artifact = f"gemma_hidden_state_{index}" + self._artifacts[artifact] = _save_tensor(self._output_dir, artifact, hidden_state) + return raw_outputs + + def record_process_hidden_states( + processor: Any, + hidden_states: tuple[torch.Tensor, ...], + attention_mask: torch.Tensor, + padding_side: str = "left", + ) -> Any: + video_features, audio_features = processor.feature_extractor(hidden_states, attention_mask, padding_side) + self._artifacts["video_features"] = _save_tensor(self._output_dir, "video_features", video_features) + self._artifacts["audio_features"] = _save_tensor(self._output_dir, "audio_features", audio_features) + return original_process_hidden_states(processor, hidden_states, attention_mask, padding_side) + + LTXGemmaTextEncoder.encode = record_encode + EmbeddingsProcessor.process_hidden_states = record_process_hidden_states + try: + outputs = self._encoder(*args, **kwargs) + finally: + LTXGemmaTextEncoder.encode = original_encode + EmbeddingsProcessor.process_hidden_states = original_process_hidden_states + if len(outputs) != 1: + raise ValueError(f"Golden capture expects one prompt, got {len(outputs)}") + output = outputs[0] + self._artifacts["video_context"] = _save_tensor(self._output_dir, "video_context", output.video_encoding) + if output.audio_encoding is not None: + self._artifacts["audio_context"] = _save_tensor(self._output_dir, "audio_context", output.audio_encoding) + self._artifacts["context_attention_mask"] = _save_tensor( + self._output_dir, "context_attention_mask", output.attention_mask + ) + return outputs + + def __getattr__(self, name: str) -> Any: + return getattr(self._encoder, name) + + +def _upstream_commit(upstream_root: Path) -> str | None: + try: + return subprocess.check_output(["git", "-C", str(upstream_root), "rev-parse", "HEAD"], text=True).strip() + except (OSError, subprocess.CalledProcessError): + return None + + +def main() -> None: + """Run one upstream T2V case and persist uncompressed reference artifacts.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", type=Path, required=True) + parser.add_argument("--upstream-root", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--prompt", required=True) + parser.add_argument("--seed", type=int, required=True) + parser.add_argument("--height", type=int, required=True) + parser.add_argument("--width", type=int, required=True) + parser.add_argument("--num-frames", type=int, required=True) + parser.add_argument("--frame-rate", type=float, default=24.0) + parser.add_argument("--video-vae", choices=("diff", "conv"), default="diff") + parser.add_argument("--image", type=Path) + parser.add_argument("--image-frame-index", type=int, default=0) + parser.add_argument("--image-strength", type=float, default=1.0) + parser.add_argument("--offload", choices=("cpu", "none"), default="cpu") + parser.add_argument( + "--attention-backend", + choices=("automatic", "pytorch"), + default="automatic", + help="Pin upstream transformer attention for a matched fidelity capture.", + ) + parser.add_argument("--capture-stage2-step2-blocks", action="store_true") + parser.add_argument( + "--diffvae-optimization", + choices=("chunked_eager", "chunked_compile"), + default="chunked_eager", + help="Select the upstream DiffVAE decoder recipe for a matched capture.", + ) + parser.add_argument( + "--diffvae-natten-backend", + choices=("automatic", "cutlass-fna"), + default="automatic", + help="Pin upstream DiffVAE NATTEN for a matched capture.", + ) + parser.add_argument( + "--deterministic-audio", + action="store_true", + help="Use deterministic CUDA kernels only while capturing decoded audio.", + ) + args = parser.parse_args() + + # Upstream imports remain local to this entry point so ordinary TeleFuser imports never depend on it. + import ltx_pipelines.utils.blocks as pipeline_blocks # type: ignore[import-not-found] + from ltx_core.loader.module_ops import ModuleOps # type: ignore[import-not-found] + from ltx_core.model.video_vae.transformer.compiling import ( + configure_natten_backend, # type: ignore[import-not-found] + ) + from ltx_core.model.video_vae.transformer.config import DiffVAEMode # type: ignore[import-not-found] + from ltx_pipelines.distilled import DistilledPipeline # type: ignore[import-not-found] + from ltx_pipelines.utils.args import ImageConditioningInput # type: ignore[import-not-found] + from ltx_pipelines.utils.media_io.encode import encode_video # type: ignore[import-not-found] + from ltx_pipelines.utils.model_paths import ModelPaths # type: ignore[import-not-found] + from ltx_pipelines.utils.types import OffloadMode # type: ignore[import-not-found] + + args.output_dir.mkdir(parents=True, exist_ok=True) + if args.image is not None and not args.image.is_file(): + raise FileNotFoundError(f"LTX-2.5 conditioning image does not exist: {args.image}") + artifacts: dict[str, Any] = {} + trajectories: dict[str, list[dict[str, Any]]] = {} + paths = ModelPaths.from_split( + transformer_path=str(args.model_root / "diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors"), + text_encoder_path=str(args.model_root / "text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors"), + video_vae_path=str( + args.model_root + / "vae" + / ( + "ltx-2.5-video-vae-bf16.safetensors" + if args.video_vae == "diff" + else "ltx-2.5-video-vae-conv-bf16.safetensors" + ) + ), + audio_vae_path=str(args.model_root / "vae/ltx-2.5-audio-vae-bf16.safetensors"), + duration_head_path=str(args.model_root / "model_patches/ltx-2.5-duration-head-bf16.safetensors"), + ) + pipeline = DistilledPipeline( + model_paths=paths, + spatial_upsampler_path=str( + args.model_root / "latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors" + ), + loras=(), + device=torch.device("cuda"), + offload_mode=OffloadMode(args.offload), + diffvae_optimization=DiffVAEMode(args.diffvae_optimization), + ) + if args.attention_backend == "pytorch": + from ltx_core.model.transformer.attention import AttentionFunction # type: ignore[import-not-found] + + pipeline.stage = pipeline.stage.with_attention(AttentionFunction.PYTORCH) + prompt_recorder = _PromptRecorder(pipeline.prompt_encoder, args.output_dir, artifacts) + pipeline.prompt_encoder = prompt_recorder + pipeline.stage = _StageRecorder( + pipeline.stage, + args.output_dir, + artifacts, + trajectories, + args.capture_stage2_step2_blocks, + ) + if args.diffvae_natten_backend == "cutlass-fna": + decoder_builder = pipeline.video_decoder._decoder_builder + + def pin_cutlass_fna(model: torch.nn.Module) -> torch.nn.Module: + configure_natten_backend(model, "cutlass-fna") + return model + + pipeline.video_decoder._decoder_builder = decoder_builder.with_module_ops( + (*decoder_builder.module_ops, ModuleOps("capture_cutlass_fna", lambda _model: True, pin_cutlass_fna)) + ) + resolved_diffvae_attention_backends: list[str | None] = ["cutlass-fna"] + else: + resolved_diffvae_attention_backends = [None] + pipeline.upsampler = _UpsamplerRecorder(pipeline.upsampler, args.output_dir, artifacts) + pipeline.video_decoder = _VideoDecoderRecorder(pipeline.video_decoder, args.output_dir, artifacts) + original_decode_audio = pipeline_blocks.vae_decode_audio + + def decode_audio(*decode_args: Any, **decode_kwargs: Any) -> Any: + with deterministic_audio_kernels(args.deterministic_audio): + return original_decode_audio(*decode_args, **decode_kwargs) + + pipeline_blocks.vae_decode_audio = decode_audio + try: + with torch.inference_mode(): + images = ( + [] + if args.image is None + else [ + ImageConditioningInput( + path=str(args.image.resolve()), + frame_idx=args.image_frame_index, + strength=args.image_strength, + ) + ] + ) + video, audio, resolved_frames, resolved_tiling = pipeline( + prompt=args.prompt, + seed=args.seed, + height=args.height, + width=args.width, + frame_rate=args.frame_rate, + images=images, + num_frames=args.num_frames, + ) + rgb = torch.cat(list(video), dim=0) + finally: + pipeline_blocks.vae_decode_audio = original_decode_audio + artifacts["decoded_rgb"] = _save_tensor(args.output_dir, "decoded_rgb", rgb) + artifacts["decoded_waveform"] = _save_tensor(args.output_dir, "decoded_waveform", audio.waveform) + mp4_path = args.output_dir / "decoded.mp4" + encode_video(rgb, int(args.frame_rate), audio, str(mp4_path), video_chunks_number=1) + components = inspect_model_pack(args.model_root, include_sha256=True) + manifest = { + "upstream_commit": _upstream_commit(args.upstream_root), + "runtime": { + "python": sys.version, + "platform": platform.platform(), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "transformers": __import__("transformers").__version__, + "gpu": torch.cuda.get_device_name(0), + "gpu_count": torch.cuda.device_count(), + "natten": _optional_package_version("natten"), + "attention_backend": args.attention_backend, + "deterministic_audio": args.deterministic_audio, + "compile": False, + "tiling": "auto", + "resolved_video_tiling": None if resolved_tiling is None else asdict(resolved_tiling), + "diffvae_optimization": args.diffvae_optimization, + "resolved_diffvae_attention_backends": resolved_diffvae_attention_backends, + }, + "request": { + "prompt": args.prompt, + "seed": args.seed, + "height": args.height, + "width": args.width, + "num_frames": args.num_frames, + "resolved_num_frames": resolved_frames, + "frame_rate": args.frame_rate, + "video_vae": args.video_vae, + "offload": args.offload, + "dtype": "bfloat16", + "prompt_normalization": prompt_recorder.normalized_prompts, + "image": ( + None + if args.image is None + else { + "path": str(args.image.resolve()), + "sha256": _sha256(args.image), + "frame_index": args.image_frame_index, + "strength": args.image_strength, + } + ), + }, + "checkpoints": {name: value.as_dict() for name, value in components.items()}, + "audio": {"sample_rate": audio.sampling_rate}, + "container": mp4_container_metadata(mp4_path), + "artifacts": artifacts, + "trajectories": trajectories, + } + manifest_path = args.output_dir / "capture_manifest.json" + manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True, ensure_ascii=True) + "\n", encoding="utf-8") + print(json.dumps({"manifest": str(manifest_path), "artifacts": sorted(artifacts)}, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/compare_ltx25_artifacts.py b/tools/validation/compare_ltx25_artifacts.py new file mode 100644 index 0000000..4946f77 --- /dev/null +++ b/tools/validation/compare_ltx25_artifacts.py @@ -0,0 +1,248 @@ +"""Compare an LTX-2.5 TeleFuser capture against an upstream golden manifest.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any + +import torch + +_EXACT_ARTIFACT_PARTS = ("token", "mask", "noise") +_DEFAULT_COSINE = 0.9999 +_DEFAULT_NRMSE = 1e-3 +_RGB_PSNR_MIN = 40.0 +_RGB_SSIM_MIN = 0.99 +_WAVEFORM_SI_SDR_MIN = 40.0 +_UPSTREAM_DIAGNOSTIC_ARTIFACTS = {"audio_features", "gemma_attention_mask", "gemma_token_ids", "video_features"} + + +def _load_manifest(root: Path) -> dict[str, Any]: + path = root / "capture_manifest.json" + if not path.is_file(): + raise FileNotFoundError(f"LTX-2.5 capture manifest does not exist: {path}") + return json.loads(path.read_text(encoding="utf-8")) + + +def _load_tensor(root: Path, descriptor: dict[str, Any]) -> torch.Tensor: + path = root / descriptor["path"] + return torch.load(path, map_location="cpu", weights_only=True) + + +def _metrics(reference: torch.Tensor, candidate: torch.Tensor) -> dict[str, float | bool]: + if reference.shape != candidate.shape: + return {"shape_match": False} + reference_float = reference.double().reshape(-1) + candidate_float = candidate.double().reshape(-1) + delta = candidate_float - reference_float + reference_norm = torch.linalg.vector_norm(reference_float) + cosine = ( + 1.0 + if reference_norm == 0 and torch.linalg.vector_norm(candidate_float) == 0 + else float(torch.nn.functional.cosine_similarity(reference_float, candidate_float, dim=0)) + ) + nrmse = float(torch.sqrt(torch.mean(delta.square())) / reference_float.square().mean().sqrt().clamp_min(1e-12)) + return { + "shape_match": True, + "dtype_match": reference.dtype == candidate.dtype, + "exact": bool(torch.equal(reference, candidate)), + "cosine": cosine, + "nrmse": nrmse, + "max_abs_error": float(delta.abs().max()), + } + + +def _is_exact_contract(name: str) -> bool: + return any(part in name for part in _EXACT_ARTIFACT_PARTS) + + +def _is_upstream_diagnostic_artifact(name: str) -> bool: + """Return whether an artifact is emitted only by the upstream recorder diagnostics.""" + return name in _UPSTREAM_DIAGNOSTIC_ARTIFACTS or name.startswith("gemma_hidden_state_") + + +def _decoded_quality(name: str, reference: torch.Tensor, candidate: torch.Tensor) -> dict[str, float]: + delta = candidate.float() - reference.float() + if name == "decoded_rgb": + rmse = torch.sqrt(torch.mean(delta.square())).item() + return { + "psnr": 20.0 * math.log10(1.0 / max(rmse, 1e-12)), + "ssim": _structural_similarity(reference, candidate), + } + reference_flat = reference.double().reshape(-1) + candidate_flat = candidate.double().reshape(-1) + scale = torch.dot(candidate_flat, reference_flat) / torch.dot(reference_flat, reference_flat).clamp_min(1e-12) + target = scale * reference_flat + noise = candidate_flat - target + return {"si_sdr": 10.0 * math.log10(float(torch.dot(target, target) / torch.dot(noise, noise).clamp_min(1e-12)))} + + +def _passes_decoded_quality(name: str, quality: dict[str, float]) -> bool: + if name == "decoded_rgb": + return quality["psnr"] >= _RGB_PSNR_MIN and quality["ssim"] >= _RGB_SSIM_MIN + return quality["si_sdr"] >= _WAVEFORM_SI_SDR_MIN + + +def _structural_similarity(reference: torch.Tensor, candidate: torch.Tensor) -> float: + """Compute SSIM for decoded RGB, using local windows when image layout is available.""" + reference_float = reference.float() + candidate_float = candidate.float() + if reference_float.ndim == 4 and reference_float.shape[-1] in {1, 3, 4}: + reference_image = reference_float.permute(0, 3, 1, 2) + candidate_image = candidate_float.permute(0, 3, 1, 2) + window_size = min(11, reference_image.shape[-2], reference_image.shape[-1]) + if window_size % 2 == 0: + window_size -= 1 + if window_size >= 3: + coords = torch.arange(window_size, dtype=reference_image.dtype, device=reference_image.device) + coords = coords - (window_size - 1) / 2 + gaussian = torch.exp(-(coords.square()) / (2 * 1.5**2)) + window = (gaussian[:, None] * gaussian[None, :]) / gaussian.sum().square() + window = window.expand(reference_image.shape[1], 1, window_size, window_size) + groups = reference_image.shape[1] + mu_reference = torch.nn.functional.conv2d(reference_image, window, padding=window_size // 2, groups=groups) + mu_candidate = torch.nn.functional.conv2d(candidate_image, window, padding=window_size // 2, groups=groups) + sigma_reference = ( + torch.nn.functional.conv2d(reference_image.square(), window, padding=window_size // 2, groups=groups) + - mu_reference.square() + ) + sigma_candidate = ( + torch.nn.functional.conv2d(candidate_image.square(), window, padding=window_size // 2, groups=groups) + - mu_candidate.square() + ) + covariance = ( + torch.nn.functional.conv2d( + reference_image * candidate_image, window, padding=window_size // 2, groups=groups + ) + - mu_reference * mu_candidate + ) + c1, c2 = 0.01**2, 0.03**2 + ssim = ((2 * mu_reference * mu_candidate + c1) * (2 * covariance + c2)) / ( + (mu_reference.square() + mu_candidate.square() + c1) * (sigma_reference + sigma_candidate + c2) + ) + return float(ssim.mean()) + + reference_flat = reference_float.reshape(-1) + candidate_flat = candidate_float.reshape(-1) + mean_reference = reference_flat.mean() + mean_candidate = candidate_flat.mean() + variance_reference = reference_flat.var(unbiased=False) + variance_candidate = candidate_flat.var(unbiased=False) + covariance = ((reference_flat - mean_reference) * (candidate_flat - mean_candidate)).mean() + c1, c2 = 0.01**2, 0.03**2 + return float( + ((2 * mean_reference * mean_candidate + c1) * (2 * covariance + c2)) + / ((mean_reference.square() + mean_candidate.square() + c1) * (variance_reference + variance_candidate + c2)) + ) + + +def compare_captures( + golden_root: Path, + candidate_root: Path, + *, + cosine_threshold: float = _DEFAULT_COSINE, + nrmse_threshold: float = _DEFAULT_NRMSE, +) -> dict[str, Any]: + """Compare shared artifact tensors and validate frozen request/checkpoint contracts.""" + golden = _load_manifest(golden_root) + candidate = _load_manifest(candidate_root) + golden_artifacts = golden.get("artifacts", {}) + candidate_artifacts = candidate.get("artifacts", {}) + if not isinstance(golden_artifacts, dict) or not isinstance(candidate_artifacts, dict): + raise ValueError("LTX-2.5 manifests must contain an artifacts object") + + results: dict[str, Any] = {} + failures: list[str] = [] + for name in sorted(set(golden_artifacts) & set(candidate_artifacts)): + metrics = _metrics( + _load_tensor(golden_root, golden_artifacts[name]), _load_tensor(candidate_root, candidate_artifacts[name]) + ) + exact_contract = _is_exact_contract(name) + decoded = name in {"decoded_rgb", "decoded_waveform"} + passed = bool(metrics.get("shape_match")) + if exact_contract: + passed = passed and bool(metrics.get("exact")) + elif decoded and passed: + quality = _decoded_quality( + name, + _load_tensor(golden_root, golden_artifacts[name]), + _load_tensor(candidate_root, candidate_artifacts[name]), + ) + metrics.update(quality) + passed = _passes_decoded_quality(name, quality) + elif passed: + passed = bool(metrics["exact"]) or bool( + metrics["cosine"] >= cosine_threshold and metrics["nrmse"] <= nrmse_threshold + ) + metrics["exact_contract"] = exact_contract + metrics["passed"] = passed + results[name] = metrics + if not passed: + failures.append(name) + + golden_request = golden.get("request", {}) + candidate_request = candidate.get("request", {}) + request_match = golden_request == candidate_request + if not request_match: + failures.append("request") + checkpoint_match = golden.get("checkpoints", {}) == candidate.get("checkpoints", {}) + if not checkpoint_match: + failures.append("checkpoints") + audio_match = golden.get("audio", {}) == candidate.get("audio", {}) + if not audio_match: + failures.append("audio") + missing_from_candidate = sorted(set(golden_artifacts) - set(candidate_artifacts)) + missing_required_from_candidate = [ + name for name in missing_from_candidate if not _is_upstream_diagnostic_artifact(name) + ] + unexpected_in_candidate = sorted(set(candidate_artifacts) - set(golden_artifacts)) + artifact_set_match = set(golden_artifacts) == set(candidate_artifacts) + artifact_contract_match = not missing_required_from_candidate and not unexpected_in_candidate + if not artifact_contract_match: + failures.append("artifact_set") + return { + "golden": str(golden_root), + "candidate": str(candidate_root), + "shared_artifacts": sorted(results), + "missing_from_candidate": missing_from_candidate, + "missing_required_from_candidate": missing_required_from_candidate, + "unexpected_in_candidate": unexpected_in_candidate, + "request_match": request_match, + "checkpoint_match": checkpoint_match, + "audio_match": audio_match, + "artifact_set_match": artifact_set_match, + "artifact_contract_match": artifact_contract_match, + "tensors": results, + "passed": not failures, + "failures": failures, + } + + +def main() -> None: + """Compare two LTX-2.5 capture directories and emit a JSON report.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("golden", type=Path) + parser.add_argument("candidate", type=Path) + parser.add_argument("--cosine-threshold", type=float, default=_DEFAULT_COSINE) + parser.add_argument("--nrmse-threshold", type=float, default=_DEFAULT_NRMSE) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + report = compare_captures( + args.golden, + args.candidate, + cosine_threshold=args.cosine_threshold, + nrmse_threshold=args.nrmse_threshold, + ) + rendered = json.dumps(report, indent=2, sort_keys=True) + "\n" + if args.output is None: + print(rendered, end="") + else: + args.output.write_text(rendered, encoding="utf-8") + if not report["passed"]: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/compare_ltx25_benchmarks.py b/tools/validation/compare_ltx25_benchmarks.py new file mode 100644 index 0000000..3655170 --- /dev/null +++ b/tools/validation/compare_ltx25_benchmarks.py @@ -0,0 +1,127 @@ +"""Decide an LTX-2.5 performance gate from matched upstream and TeleFuser reports.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +_NOISE_BAND = 0.02 +_RUNTIME_FIELDS = ("torch_version", "cuda_version", "natten_version", "natten_has_libnatten", "gpus") + + +def _load_report(path: Path) -> dict[str, Any]: + """Load one benchmark report and validate that it is a JSON object.""" + report = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(report, dict): + raise ValueError(f"LTX-2.5 benchmark report must be a JSON object: {path}") + return report + + +def _summary(report: dict[str, Any], *path: str) -> dict[str, Any]: + value: Any = report + for key in path: + if not isinstance(value, dict) or key not in value: + raise ValueError(f"LTX-2.5 benchmark report is missing {'/'.join(path)}") + value = value[key] + if not isinstance(value, dict): + raise ValueError(f"LTX-2.5 benchmark summary must be an object: {'/'.join(path)}") + if not isinstance(value.get("p50_seconds"), (int, float)) or not isinstance(value.get("count"), int): + raise ValueError(f"LTX-2.5 benchmark summary is malformed: {'/'.join(path)}") + return value + + +def _compare_measurement(candidate: float, upstream: float) -> dict[str, float | str]: + if upstream <= 0: + raise ValueError("upstream p50 must be positive") + relative_delta = candidate / upstream - 1.0 + if relative_delta > _NOISE_BAND: + status = "failed" + elif relative_delta < -_NOISE_BAND: + status = "passed" + else: + status = "inconclusive" + return { + "upstream_p50_seconds": upstream, + "telefuser_p50_seconds": candidate, + "relative_delta": relative_delta, + "status": status, + } + + +def compare_benchmarks( + upstream: dict[str, Any], + candidate: dict[str, Any], + *, + minimum_samples: int = 5, +) -> dict[str, Any]: + """Apply the frozen request/runtime and 2%-noise performance contracts.""" + if minimum_samples < 1: + raise ValueError("minimum_samples must be positive") + if upstream.get("implementation") != "upstream": + raise ValueError("upstream report implementation must be 'upstream'") + if candidate.get("implementation") != "telefuser": + raise ValueError("candidate report implementation must be 'telefuser'") + + request_match = upstream.get("request") == candidate.get("request") + upstream_runtime = upstream.get("runtime", {}) + candidate_runtime = candidate.get("runtime", {}) + if not isinstance(upstream_runtime, dict) or not isinstance(candidate_runtime, dict): + raise ValueError("benchmark reports must contain runtime objects") + runtime_match = {field: upstream_runtime.get(field) == candidate_runtime.get(field) for field in _RUNTIME_FIELDS} + + phases = { + "cold_end_to_end": (_summary(upstream, "cold", "end_to_end"), _summary(candidate, "cold", "end_to_end")), + "warm_end_to_end": (_summary(upstream, "warm"), _summary(candidate, "warm")), + } + sample_counts = { + name: {"upstream": upstream_summary["count"], "telefuser": candidate_summary["count"]} + for name, (upstream_summary, candidate_summary) in phases.items() + } + sufficient_samples = all(count >= minimum_samples for counts in sample_counts.values() for count in counts.values()) + measurements = { + name: _compare_measurement(candidate_summary["p50_seconds"], upstream_summary["p50_seconds"]) + for name, (upstream_summary, candidate_summary) in phases.items() + } + statuses = [measurement["status"] for measurement in measurements.values()] + passed = ( + request_match + and all(runtime_match.values()) + and sufficient_samples + and all(status == "passed" for status in statuses) + ) + return { + "request_match": request_match, + "runtime_match": runtime_match, + "minimum_samples": minimum_samples, + "sample_counts": sample_counts, + "sufficient_samples": sufficient_samples, + "noise_band": _NOISE_BAND, + "measurements": measurements, + "passed": passed, + } + + +def main() -> None: + """Compare matched raw LTX-2.5 benchmark reports and write a decision artifact.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("upstream", type=Path) + parser.add_argument("telefuser", type=Path) + parser.add_argument("--minimum-samples", type=int, default=5) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + report = compare_benchmarks( + _load_report(args.upstream), _load_report(args.telefuser), minimum_samples=args.minimum_samples + ) + rendered = json.dumps(report, indent=2, sort_keys=True) + "\n" + if args.output is None: + print(rendered, end="") + else: + args.output.write_text(rendered, encoding="utf-8") + if not report["passed"]: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/inspect_ltx25_checkpoints.py b/tools/validation/inspect_ltx25_checkpoints.py new file mode 100644 index 0000000..24d88cf --- /dev/null +++ b/tools/validation/inspect_ltx25_checkpoints.py @@ -0,0 +1,36 @@ +"""Write a metadata-only provenance manifest for an LTX-2.5 distilled model pack.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +from telefuser.models.ltx25.checkpoint import inspect_model_pack + + +def main() -> None: + """Inspect required LTX-2.5 split checkpoints without materializing tensors.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--sha256", + action="store_true", + help="Hash checkpoint payloads; this reads every file in full.", + ) + args = parser.parse_args() + + components = inspect_model_pack(args.model_root, include_sha256=args.sha256) + manifest: dict[str, object] = { + "model_root": str(args.model_root.resolve()), + "sha256_included": args.sha256, + "components": {name: metadata.as_dict() for name, metadata in components.items()}, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(manifest, indent=2, sort_keys=True, ensure_ascii=True) + "\n", encoding="utf-8") + print(json.dumps({"output": str(args.output), "components": sorted(manifest["components"])}, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/ltx25_capture_utils.py b/tools/validation/ltx25_capture_utils.py new file mode 100644 index 0000000..3225fa4 --- /dev/null +++ b/tools/validation/ltx25_capture_utils.py @@ -0,0 +1,86 @@ +"""Shared runtime controls for LTX-2.5 validation captures.""" + +from __future__ import annotations + +import hashlib +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Iterator + +import torch + + +@contextmanager +def deterministic_audio_kernels(enabled: bool) -> Iterator[None]: + """Stabilize CUDA convolution selection while capturing decoded audio.""" + if not enabled: + yield + return + benchmark = torch.backends.cudnn.benchmark + deterministic = torch.backends.cudnn.deterministic + deterministic_algorithms = torch.are_deterministic_algorithms_enabled() + try: + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + torch.use_deterministic_algorithms(True) + yield + finally: + torch.use_deterministic_algorithms(deterministic_algorithms) + torch.backends.cudnn.deterministic = deterministic + torch.backends.cudnn.benchmark = benchmark + + +def mp4_container_metadata(path: Path) -> dict[str, Any]: + """Return JSON-safe provenance for an encoded MP4 container and its streams.""" + import av + + def optional_string(value: object | None) -> str | None: + return None if value is None else str(value) + + def stream_metadata(stream: Any) -> dict[str, Any]: + codec = stream.codec_context + result: dict[str, Any] = { + "type": stream.type, + "codec": codec.name, + "time_base": optional_string(stream.time_base), + "duration": stream.duration, + "frames": stream.frames, + "bit_rate": codec.bit_rate, + "metadata": dict(stream.metadata), + } + if stream.type == "video": + result.update( + { + "width": codec.width, + "height": codec.height, + "pixel_format": optional_string(codec.format), + "frame_rate": optional_string(stream.average_rate), + } + ) + elif stream.type == "audio": + result.update( + { + "sample_rate": codec.sample_rate, + "channels": codec.channels, + "layout": optional_string(codec.layout), + "sample_format": optional_string(codec.format), + } + ) + return result + + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + with av.open(str(path)) as container: + return { + "path": path.name, + "sha256": digest.hexdigest(), + "size_bytes": path.stat().st_size, + "format": container.format.name, + "duration": container.duration, + "start_time": container.start_time, + "bit_rate": container.bit_rate, + "metadata": dict(container.metadata), + "streams": [stream_metadata(stream) for stream in container.streams], + }