diff --git a/README.md b/README.md index 0d1c9e83c..97b3913fd 100644 --- a/README.md +++ b/README.md @@ -41,6 +41,8 @@ See also: > Currently, the development personnel of this project are limited, with most of the work handled by [Artiprocher](https://github.com/Artiprocher) and [mi804](https://github.com/mi804). Therefore, the progress of new feature development will be relatively slow, and the speed of responding to and resolving issues is limited. We apologize for this and ask developers to understand. +- **September 10, 2026** We have integrated [LTX-2.5](https://modelscope.cn/models/Lightricks/LTX-2.5), the latest audio-video generation model from Lightricks. The features include text-to-audio/video with automatic duration prediction, image-to-audio/video with keyframe interpolation, audio-to-video, audio-video retake, IC-LoRA pixel spatial upscaling, text-to-audio, INT8 quantized inference, low VRAM inference, and training. For details, please refer to the [documentation](/docs/en/Model_Details/LTX-2.md) and [code](/examples/ltx2/). + - **September 10, 2026** We have open-sourced [DiffSynth-ComfyUI](https://github.com/modelscope/DiffSynth-ComfyUI). Now you can run model inference with DiffSynth-Studio in [ComfyUI](https://github.com/Comfy-Org/ComfyUI). - **September 8, 2026** We trained a [MiniMax-H3 Training Adapter](https://www.modelscope.cn/models/DiffSynth-Studio/MiniMax-H3-TrainingAdapter) for the CFG-distilled MiniMax-H3 base, a DeCFG LoRA training adapter offered in FL2VA and Ref2VA versions. The training data comes from the self-generated [MiniMax-H3-Self-Generated-Dataset](https://www.modelscope.cn/datasets/DiffSynth-Studio/MiniMax-H3-Self-Generated-Dataset). Based on this adapter, we trained two toy models, [MiniMax-H3-Songyu-LoRA](https://www.modelscope.cn/models/mibei0804/MiniMax-H3-Songyu-LoRA) (character identity, FL2VA) and [MiniMax-H3-Ref2VA-FirstFrame-Lineart](https://www.modelscope.cn/models/mibei0804/MiniMax-H3-Ref2VA-FirstFrame-Lineart) (lineart first-frame control, Ref2VA), as references for inference and fine-tuning. @@ -642,6 +644,16 @@ https://github.com/Artiprocher/DiffSynth-Studio/assets/35051019/59fb2f7b-8de0-44 | JoyAI-Image | [jd-opensource/JoyAI-Image-Edit](https://modelscope.cn/models/jd-opensource/JoyAI-Image-Edit) | [code](/examples/joyai_image/model_inference/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_inference_low_vram/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_training/full/JoyAI-Image-Edit.sh) | [code](/examples/joyai_image/model_training/validate_full/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_training/lora/JoyAI-Image-Edit.sh) | [code](/examples/joyai_image/model_training/validate_lora/JoyAI-Image-Edit.py) | | ERNIE-Image | [PaddlePaddle/ERNIE-Image](https://www.modelscope.cn/models/PaddlePaddle/ERNIE-Image) | [code](/examples/ernie_image/model_inference/ERNIE-Image.py) | [code](/examples/ernie_image/model_inference_low_vram/ERNIE-Image.py) | [code](/examples/ernie_image/model_training/full/ERNIE-Image.sh) | [code](/examples/ernie_image/model_training/validate_full/ERNIE-Image.py) | [code](/examples/ernie_image/model_training/lora/ERNIE-Image.sh) | [code](/examples/ernie_image/model_training/validate_lora/ERNIE-Image.py) | | ERNIE-Image | [PaddlePaddle/ERNIE-Image-Turbo](https://www.modelscope.cn/models/PaddlePaddle/ERNIE-Image-Turbo) | [code](/examples/ernie_image/model_inference/ERNIE-Image-Turbo.py) | [code](/examples/ernie_image/model_inference_low_vram/ERNIE-Image-Turbo.py) | — | — | — | — | +| LTX-2.5 | [Lightricks/LTX-2.5: OneStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py) | [code](/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py) | [code](/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py) | +| LTX-2.5 | [Lightricks/LTX-2.5: TwoStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: OneStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: TwoStagePipeline-A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: TwoStagePipeline-Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2A.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler) | [code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: INT8-ConvRot](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py) | - | - | - | - | | LTX-2 | [jd-opensource/JoyAI-Echo](https://modelscope.cn/models/jd-opensource/JoyAI-Echo) | [code](/examples/ltx2/model_inference/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_inference_low_vram/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_training/full/JoyAI-Echo-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_training/lora/JoyAI-Echo-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/JoyAI-Echo-T2AV.py) | | LTX-2 | [Lightricks/LTX-2.3: OneStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-I2AV-OneStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-OneStage.py) | [code](/examples/ltx2/model_training/full/LTX-2.3-I2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/LTX-2.3-I2AV.py) | [code](/examples/ltx2/model_training/lora/LTX-2.3-I2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2.3-I2AV.py) | | LTX-2 | [Lightricks/LTX-2.3: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-I2AV-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-TwoStage.py) | - | - | - | - | diff --git a/README_zh.md b/README_zh.md index 51fb0244e..81baa8f7f 100644 --- a/README_zh.md +++ b/README_zh.md @@ -41,6 +41,8 @@ > 目前本项目的开发人员有限，大部分工作由 [Artiprocher](https://github.com/Artiprocher) 和 [mi804](https://github.com/mi804) 负责，因此新功能的开发进展会比较缓慢，issue 的回复和解决速度有限，我们对此感到非常抱歉，请各位开发者理解。 +- **2026年9月10日** 我们接入了 [LTX-2.5](https://modelscope.cn/models/Lightricks/LTX-2.5)，这是 Lightricks 最新的音视频联合生成模型。支持的功能包括自动时长预测的文生音视频、关键帧插值的图生音视频、音频驱动视频、音视频区域重生成、IC-LoRA 像素空间上采样、文生音频、INT8 量化推理、低显存推理以及模型训练。详情请参考[文档](/docs/zh/Model_Details/LTX-2.md)和[示例代码](/examples/ltx2/)。 + - **2026年9月10日** 我们开源了 [DiffSynth-ComfyUI](https://github.com/modelscope/DiffSynth-ComfyUI)，现在，你可以在 [ComfyUI](https://github.com/Comfy-Org/ComfyUI) 中使用 DiffSynth-Studio 运行模型推理了。 - **2026年9月8日** 我们为 CFG 蒸馏的 MiniMax-H3 基座训练了 [MiniMax-H3 Training Adapter](https://www.modelscope.cn/models/DiffSynth-Studio/MiniMax-H3-TrainingAdapter)，这是一个 DeCFG LoRA 训练 Adapter，提供 FL2VA 与 Ref2VA 两个版本。训练数据来自自生成数据集 [MiniMax-H3-Self-Generated-Dataset](https://www.modelscope.cn/datasets/DiffSynth-Studio/MiniMax-H3-Self-Generated-Dataset)。基于该 Adapter，我们训练了两个 toy 模型，[MiniMax-H3-Songyu-LoRA](https://www.modelscope.cn/models/mibei0804/MiniMax-H3-Songyu-LoRA)（角色身份，FL2VA）与 [MiniMax-H3-Ref2VA-FirstFrame-Lineart](https://www.modelscope.cn/models/mibei0804/MiniMax-H3-Ref2VA-FirstFrame-Lineart)（线稿首帧控制，Ref2VA），供推理与微调参考。 @@ -603,6 +605,13 @@ https://github.com/Artiprocher/DiffSynth-Studio/assets/35051019/59fb2f7b-8de0-44 | JoyAI-Image | [jd-opensource/JoyAI-Image-Edit](https://modelscope.cn/models/jd-opensource/JoyAI-Image-Edit) | [code](/examples/joyai_image/model_inference/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_inference_low_vram/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_training/full/JoyAI-Image-Edit.sh) | [code](/examples/joyai_image/model_training/validate_full/JoyAI-Image-Edit.py) | [code](/examples/joyai_image/model_training/lora/JoyAI-Image-Edit.sh) | [code](/examples/joyai_image/model_training/validate_lora/JoyAI-Image-Edit.py) | | ERNIE-Image | [PaddlePaddle/ERNIE-Image](https://www.modelscope.cn/models/PaddlePaddle/ERNIE-Image) | [code](/examples/ernie_image/model_inference/ERNIE-Image.py) | [code](/examples/ernie_image/model_inference_low_vram/ERNIE-Image.py) | [code](/examples/ernie_image/model_training/full/ERNIE-Image.sh) | [code](/examples/ernie_image/model_training/validate_full/ERNIE-Image.py) | [code](/examples/ernie_image/model_training/lora/ERNIE-Image.sh) | [code](/examples/ernie_image/model_training/validate_lora/ERNIE-Image.py) | | ERNIE-Image | [PaddlePaddle/ERNIE-Image-Turbo](https://www.modelscope.cn/models/PaddlePaddle/ERNIE-Image-Turbo) | [code](/examples/ernie_image/model_inference/ERNIE-Image-Turbo.py) | [code](/examples/ernie_image/model_inference_low_vram/ERNIE-Image-Turbo.py) | — | — | — | — | +| LTX-2.5 | [Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py) | [code](/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py) | +| LTX-2.5 | [Lightricks/LTX-2.5: DistilledPipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-I2AV-DistilledPipeline.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-DistilledPipeline.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: TwoStagePipeline-A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: TwoStagePipeline-Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler) | [code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2A.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py) | - | - | - | - | +| LTX-2.5 | [Lightricks/LTX-2.5: INT8-ConvRot](https://www.modelscope.cn/models/Lightricks/LTX-2.5) | [code](/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py) | - | - | - | - | | LTX-2 | [jd-opensource/JoyAI-Echo](https://modelscope.cn/models/jd-opensource/JoyAI-Echo) | [code](/examples/ltx2/model_inference/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_inference_low_vram/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_training/full/JoyAI-Echo-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/JoyAI-Echo-T2AV.py) | [code](/examples/ltx2/model_training/lora/JoyAI-Echo-T2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/JoyAI-Echo-T2AV.py) | | LTX-2 | [Lightricks/LTX-2.3: OneStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-I2AV-OneStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-OneStage.py) | [code](/examples/ltx2/model_training/full/LTX-2.3-I2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_full/LTX-2.3-I2AV.py) | [code](/examples/ltx2/model_training/lora/LTX-2.3-I2AV-splited.sh) | [code](/examples/ltx2/model_training/validate_lora/LTX-2.3-I2AV.py) | | LTX-2 | [Lightricks/LTX-2.3: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.3) | [code](/examples/ltx2/model_inference/LTX-2.3-I2AV-TwoStage.py) | [code](/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-TwoStage.py) | - | - | - | - | diff --git a/diffsynth/configs/model_configs.py b/diffsynth/configs/model_configs.py index d9b3cff49..ff69a5d3a 100644 --- a/diffsynth/configs/model_configs.py +++ b/diffsynth/configs/model_configs.py @@ -850,7 +850,7 @@ "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_text_encoder.LTX2TextEncoderPostModulesStateDictConverter", }, { - # Example: ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors") + # Example: ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors") "model_hash": "aed408774d694a2452f69936c32febb5", "model_name": "ltx2_latent_upsampler", "model_class": "diffsynth.models.ltx2_upsampler.LTX2LatentUpsampler", @@ -895,6 +895,104 @@ "extra_kwargs": {"separated_audio_video": True, "embedding_dim_gemma": 3840, "num_layers_gemma": 49, "video_attention_heads": 32, "video_attention_head_dim": 128, "audio_attention_heads": 32, "audio_attention_head_dim": 64, "num_connector_layers": 8, "apply_gated_attention": True}, "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_text_encoder.LTX2TextEncoderPostModulesStateDictConverter", }, + { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors") + "model_hash": "7960c5dc4626650824e36f65a8e992e9", + "model_name": "ltx2_dit", + "model_class": "diffsynth.models.ltx2_dit.LTXModel", + "extra_kwargs": {"caption_channels": None, "apply_gated_attention": True, "cross_attention_adaln": True, "ff_bias": False, "use_keyframes_abs_pos_embedding": True, "use_tokenwise_av_ca_scale_shift": True}, + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_dit.LTXModelStateDictConverter", + }, + { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-comfy-int8-convrot.safetensors") + "model_hash": "57343d320cac0bbba58a488b8ebe7187", + "model_name": "ltx2_dit", + "model_class": "diffsynth.models.ltx2_dit.LTXModel", + "extra_kwargs": {"caption_channels": None, "apply_gated_attention": True, "cross_attention_adaln": True, "ff_bias": False, "use_keyframes_abs_pos_embedding": True, "use_tokenwise_av_ca_scale_shift": True}, + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_dit.LTXModelStateDictConverter", + "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["timestep_embedder.linear_1", "timestep_embedder.linear_2", "adaln_single.linear", "audio_adaln_single.linear", "prompt_adaln_single.linear", "audio_prompt_adaln_single.linear", "av_ca_a2v_gate_adaln_single.linear", "av_ca_audio_scale_shift_adaln_single.linear", "av_ca_v2a_gate_adaln_single.linear", "av_ca_video_scale_shift_adaln_single.linear", "patchify_proj", "audio_patchify_proj", "proj_out", "audio_proj_out", "to_gate_logits"]}, + }, + { + # Example: ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors") + "model_hash": "8f3c146ff3d584392236c5b29d26146c", + "model_name": "ltx2_text_encoder_post_modules", + "model_class": "diffsynth.models.ltx25_text_encoder.LTX25TextEncoderPostModules", + }, + { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors") + "model_hash": "055700dc619165899bebb5162f699cd2", + "model_name": "ltx2_text_encoder", + "model_class": "diffsynth.models.ltx25_text_encoder.LTX25TextEncoder", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25TextEncoderStateDictConverter", + }, + { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors") + "model_hash": "4743ded7a5725b6589bccdb62512723b", + "model_name": "ltx2_text_encoder", + "model_class": "diffsynth.models.ltx25_text_encoder.LTX25TextEncoder", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_text_encoder.LTX25TextEncoderStateDictConverter", + "quant_config": {"method": "comfy_kitchen_int8_w8a8", "load_prequantized": True, "exclude_modules": ["lm_head", "embedding_projection", "patch_dense"]}, + }, + { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors") + "model_hash": "e19205490f01801d0a7b6d3aba61e26e", + "model_name": "ltx2_video_vae_encoder", + "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder", + "extra_kwargs": {"encoder_version": "ltx-2.3"}, + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoEncoderStateDictConverter", + }, + { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors") + "model_hash": "e19205490f01801d0a7b6d3aba61e26e", + "model_name": "ltx25_diffusion_video_vae_decoder", + "model_class": "diffsynth.models.ltx25_diffusion_video_vae.LTX25DiffusionVideoDecoder", + "extra_kwargs": {"stage_channels": [2048, 1024, 512, 512, 256], "stage_depths": [4, 6, 4, 2, 8], "stage_kernels": [[3, 7, 7], [3, 7, 7], [3, 5, 5], [3, 5, 5], [11, 11, 11]], "stage5_kernel": [11, 11, 11], "timestep_scale_multiplier": 1000.0, "default_num_inference_steps": 1, "model_output_type": "x0"}, + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_diffusion_video_vae.LTX25DiffusionVideoDecoderStateDictConverter", + }, + { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors") + "model_hash": "a1d642eecae96baa9c31d4e405564f49", + "model_name": "ltx2_video_vae_encoder", + "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder", + "extra_kwargs": {"encoder_version": "ltx-2.3"}, + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoEncoderStateDictConverter", + }, + { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors") + "model_hash": "a1d642eecae96baa9c31d4e405564f49", + "model_name": "ltx2_video_vae_decoder", + "model_class": "diffsynth.models.ltx2_video_vae.LTX2VideoDecoder", + "extra_kwargs": {"decoder_version": "ltx-2.3"}, + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_video_vae.LTX2VideoDecoderStateDictConverter", + }, + { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors") + "model_hash": "c2488315f13356abb806f9f217f1e803", + "model_name": "ltx2_audio_vae_decoder", + "model_class": "diffsynth.models.ltx2_audio_vae.LTX2AudioDecoder", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2AudioDecoderStateDictConverter", + }, + { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors") + "model_hash": "c2488315f13356abb806f9f217f1e803", + "model_name": "ltx2_audio_vocoder", + "model_class": "diffsynth.models.ltx2_audio_vae.LTX2VocoderWithBWE", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2VocoderStateDictConverter", + }, + { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors") + "model_hash": "c2488315f13356abb806f9f217f1e803", + "model_name": "ltx2_audio_vae_encoder", + "model_class": "diffsynth.models.ltx2_audio_vae.LTX2AudioEncoder", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx2_audio_vae.LTX2AudioEncoderStateDictConverter", + }, + { + # Example: ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors") + "model_hash": "35840495e440a4f00946450269299bd6", + "model_name": "ltx25_duration_head", + "model_class": "diffsynth.models.ltx25_duration_head.LTX25DurationHead", + "state_dict_converter": "diffsynth.utils.state_dict_converters.ltx25_duration_head.LTX25DurationHeadStateDictConverter", + }, ] anima_series = [ { diff --git a/diffsynth/configs/vram_management_module_maps.py b/diffsynth/configs/vram_management_module_maps.py index 18d925678..d5366be3a 100644 --- a/diffsynth/configs/vram_management_module_maps.py +++ b/diffsynth/configs/vram_management_module_maps.py @@ -273,20 +273,37 @@ "diffsynth.models.ltx2_dit.LTXModel": { "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_dit.BasicAVTransformerBlock": "diffsynth.core.vram.layers.AutoWrappedNonRecurseModule", + }, + "diffsynth.models.ltx25_text_encoder.LTX25TextEncoder": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule", + "torch.nn.LayerNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "transformers.models.gemma4_unified.modeling_gemma4_unified.Gemma4UnifiedTextDecoderLayer": "diffsynth.core.vram.layers.AutoWrappedNonRecurseModule", + "transformers.models.gemma4_unified.modeling_gemma4_unified.Gemma4UnifiedRMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "transformers.models.gemma4_unified.modeling_gemma4_unified.Gemma4UnifiedTextRotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.ltx25_text_encoder.LTX25TextEncoderPostModules": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx25_text_encoder.LTX25Embeddings1DConnector": "diffsynth.core.vram.layers.AutoWrappedModule", + }, + "diffsynth.models.ltx25_diffusion_video_vae.LTX25DiffusionVideoDecoder": { + "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear", + "torch.nn.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx25_diffusion_video_vae.DiffusionVideoDecoder": "diffsynth.core.vram.layers.AutoWrappedModule", }, "diffsynth.models.ltx2_upsampler.LTX2LatentUpsampler": { - "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", - "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule", - "torch.nn.GroupNorm": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_upsampler.LTX2LatentUpsampler": "diffsynth.core.vram.layers.AutoWrappedModule", }, "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder": { - "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_video_vae.LTX2VideoEncoder": "diffsynth.core.vram.layers.AutoWrappedModule", }, "diffsynth.models.ltx2_video_vae.LTX2VideoDecoder": { - "torch.nn.Conv3d": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_video_vae.LTX2VideoDecoder": "diffsynth.core.vram.layers.AutoWrappedModule", }, "diffsynth.models.ltx2_audio_vae.LTX2AudioDecoder": { - "torch.nn.Conv2d": "diffsynth.core.vram.layers.AutoWrappedModule", + "diffsynth.models.ltx2_audio_vae.LTX2AudioDecoder": "diffsynth.core.vram.layers.AutoWrappedModule", }, "diffsynth.models.ltx2_audio_vae.LTX2Vocoder": { "torch.nn.Conv1d": "diffsynth.core.vram.layers.AutoWrappedModule", diff --git a/diffsynth/diffusion/__init__.py b/diffsynth/diffusion/__init__.py index d285482de..f5d9ea194 100644 --- a/diffsynth/diffusion/__init__.py +++ b/diffsynth/diffusion/__init__.py @@ -1,4 +1,4 @@ -from .flow_match import FlowMatchScheduler, HiDreamO1FlashScheduler +from .flow_match import FlowMatchScheduler, HiDreamO1FlashScheduler, AncestralFlowMatchScheduler from .training_module import DiffusionTrainingModule from .logger import ModelLogger from .runner import launch_training_task, launch_data_process_task diff --git a/diffsynth/diffusion/flow_match.py b/diffsynth/diffusion/flow_match.py index 2d7381427..df9ff284e 100644 --- a/diffsynth/diffusion/flow_match.py +++ b/diffsynth/diffusion/flow_match.py @@ -451,3 +451,28 @@ def step(self, model_output, timestep, sample): noise = self.clip_noise(torch.randn(denoised.shape, device=denoised.device, dtype=denoised.dtype)) sample = sigma_ * noise * self.noise_scale_schedule[timestep_id] + (1.0 - sigma_) * denoised return sample + + +class AncestralFlowMatchScheduler(FlowMatchScheduler): + + def __init__(self, template="LTX-2", eta=1.0, s_noise=1.0, noise_seed=0, rand_device="cpu"): + super().__init__(template) + self.eta = eta + self.s_noise = s_noise + self.generator = torch.Generator(device=rand_device).manual_seed(noise_seed) + + def step(self, model_output, timestep, sample, **kwargs): + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + sigma_ = self.sigmas[timestep_id + 1] if timestep_id + 1 < len(self.timesteps) else torch.zeros_like(sigma) + denoised = sample.float() - model_output.float() * sigma.float() + if sigma_ == 0: + return denoised.to(sample.dtype) + sigma_down = sigma_ * (1.0 + (sigma_ / sigma - 1.0) * self.eta) + ratio = sigma_down / sigma + prev_sample = ratio * sample.float() + (1.0 - ratio) * denoised + alpha_next, alpha_down = 1.0 - sigma_, 1.0 - sigma_down + renoise_coeff = (sigma_ ** 2 - sigma_down ** 2 * alpha_next ** 2 / alpha_down ** 2).clamp(min=0).sqrt() + noise = torch.randn(sample.shape, generator=self.generator, dtype=sample.dtype, device=self.generator.device).to(sample.device) + prev_sample = alpha_next / alpha_down * prev_sample + noise.float() * self.s_noise * renoise_coeff + return prev_sample.to(sample.dtype) diff --git a/diffsynth/models/ltx25_diffusion_video_vae.py b/diffsynth/models/ltx25_diffusion_video_vae.py new file mode 100644 index 000000000..d340ec3a1 --- /dev/null +++ b/diffsynth/models/ltx25_diffusion_video_vae.py @@ -0,0 +1,5014 @@ +from __future__ import annotations + +import dataclasses +import itertools +import logging +import math + +from tqdm import tqdm +from collections.abc import Iterator, Sequence +from dataclasses import dataclass, replace +from enum import Enum +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Callable, Final, List, Literal, NamedTuple, Protocol, Tuple, TYPE_CHECKING + +import numpy as np +import torch +import torch.nn.functional as F +from einops import rearrange +from torch import nn +from torch.nn import functional + + +class Disposable: + pass +class VideoDecoder: + pass +def _clip_generators(count, generator): + if isinstance(generator, Sequence): + if len(generator) != count: + raise ValueError(f"decode_single_frames got {count} latents and {len(generator)} generators") + return generator + return [generator] * count + + +def iter_decoded_single_frames(decoder, latents, generator=None): + generators = _clip_generators(len(latents), generator) + for index, (latent, item_generator) in enumerate(zip(latents, generators, strict=True)): + if latent.ndim != 5 or latent.shape[2] != 1: + raise ValueError( + f"decode_single_frames expects (B, C, 1, H, W) latents, got {tuple(latent.shape)} at index {index}" + ) + chunks = list(decoder.decode_video(latent, tiling_config=None, generator=item_generator)) + if not chunks: + raise RuntimeError(f"Decoder returned no pixels for single-frame latent {index}") + yield torch.cat(chunks, dim=0) + + +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: + 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): + 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 + + +class VideoPixelShape(NamedTuple): + batch: int + frames: int + height: int + width: int + fps: float + + +class SpatioTemporalScaleFactors(NamedTuple): + 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": + 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) + +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( + batch=shape[0], + channels=shape[1], + frames=shape[2], + height=shape[3], + width=shape[4], + ) + + @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): + 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]) + + @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: + 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: + 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 + +@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 + 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, + ) + + +def rms_norm(x: torch.Tensor, weight: torch.Tensor | None = None, eps: float = 1e-6) -> torch.Tensor: + return torch.nn.functional.rms_norm(x, (x.shape[-1],), weight=weight, eps=eps) + + +def to_velocity( + sample: torch.Tensor, + sigma: float | torch.Tensor, + denoised_sample: torch.Tensor, + calc_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + 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: + if isinstance(sigma, torch.Tensor): + sigma = sigma.to(calc_dtype) + return (sample.to(calc_dtype) - velocity.to(calc_dtype) * sigma).to(sample.dtype) + + +def compute_trapezoidal_mask_1d( + length: int, + ramp_left: int, + ramp_right: int, + left_starts_from_0: bool = False, +) -> torch.Tensor: + 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: + 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: list[DimensionInterval] + + +@dataclass(frozen=True) +class LatentIntervals: + original_shape: torch.Size + dimension_intervals: tuple[DimensionIntervals, ...] + + +SplitOperation = Callable[[int], DimensionIntervals] + + +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: + 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 + + +def _grow_last_tile_to_min(intervals: list[DimensionInterval], min_tile_size: int) -> list[DimensionInterval]: + 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: + 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: + 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: + 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_by_count_temporal_causal( + num_tiles: int, overlap: int = 0, min_tile_size: int | None = None +) -> SplitOperation: + 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_at_seams(boundaries: Sequence[int], num_tiles: int, overlap: int = 0) -> SplitOperation: + boundaries = tuple(boundaries) + 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 len(boundaries) < 2 or boundaries[0] != 0: + raise ValueError(f"boundaries must start at 0 and hold at least one segment, got {list(boundaries)}") + if any(b <= a for a, b in itertools.pairwise(boundaries)): + raise ValueError(f"boundaries must be strictly increasing, got {list(boundaries)}") + n_segments = len(boundaries) - 1 + n_tiles = min(num_tiles, n_segments) + base, leftover = divmod(n_segments, n_tiles) + counts = [base + (1 if index < leftover else 0) for index in range(n_tiles)] + + def split(dim_size: int) -> DimensionIntervals: + if boundaries[-1] != dim_size - 1: + raise ValueError(f"boundaries must end at the last cell ({dim_size - 1}), got {boundaries[-1]}") + intervals: list[DimensionInterval] = [] + cursor = 0 + for tile_index, count in enumerate(counts): + resume = boundaries[cursor] + 1 + start = 0 if tile_index == 0 else max(0, resume - overlap) + cursor += count + intervals.append( + DimensionInterval( + start=start, + end=boundaries[cursor] + 1, + left_ramp=0 if tile_index == 0 else resume - start, + right_ramp=0, + ) + ) + return DimensionIntervals(intervals=intervals) + + return split + + +def split_by_count(num_tiles: int, overlap: int = 0, min_tile_size: int | None = None) -> SplitOperation: + 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 + + +def identity_mapping_operation( + intervals: DimensionIntervals, + *, + rectangular: bool = False, +) -> tuple[list[slice], list[torch.Tensor]]: + mask_1d = compute_rectangular_mask_1d if rectangular else compute_trapezoidal_mask_1d + out_slices: list[slice] = [] + masks: list[torch.Tensor] = [] + for iv in intervals.intervals: + out_slices.append(slice(iv.start, iv.end)) + masks.append(mask_1d(iv.end - iv.start, iv.left_ramp, iv.right_ramp)) + return out_slices, masks + + +class Tile(NamedTuple): + 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: + 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: + 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 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]]: + 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 + + +@dataclass(frozen=True) +class DimensionTilingConfig: + 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: + return self.num_tiles > 1 or self.overlap > 0 + +@dataclass(frozen=True) +class DimensionSizeConfig: + 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: + return self.tile_size > 0 + + +@dataclass(frozen=True) +class TileCountConfig: + frames: DimensionTilingConfig = DimensionTilingConfig() + height: DimensionTilingConfig = DimensionTilingConfig() + width: DimensionTilingConfig = DimensionTilingConfig() + + def validate(self, scale_factors: SpatioTemporalScaleFactors, video_shape: VideoPixelShape) -> None: + 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]: + 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), + ) + +@dataclass(frozen=True) +class TileSizeConfig: + frames: DimensionSizeConfig = DimensionSizeConfig() + height: DimensionSizeConfig = DimensionSizeConfig() + width: DimensionSizeConfig = DimensionSizeConfig() + + def validate(self, scale_factors: SpatioTemporalScaleFactors, video_shape: VideoPixelShape) -> None: + 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), + ) + + def to_splitters( + self, + scale_factors: SpatioTemporalScaleFactors, + min_tile_size: tuple[int, int, int] | None = None, + *, + causal_temporal: bool = True, + ) -> tuple[SplitOperation, SplitOperation, SplitOperation]: + 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), + ) + +TilingConfig = TileSizeConfig | TileCountConfig + + + + +def _assert_video_on_vae_grid( + scale_factors: SpatioTemporalScaleFactors, + video_shape: VideoPixelShape, +) -> bool: + 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: + 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: + 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: + 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}.") + + +class DiffVAEMode(Enum): + COMBINED_COMPILE = "combined_compile" + CHUNKED_COMPILE = "chunked_compile" + CHUNKED_EAGER = "chunked_eager" + BLACKWELL_DSL = "blackwell_dsl" + + def resolve(self): + return self + + +class NAttentionKind(Enum): + TRITON = "triton" + EAGER_SDPA = "eager_sdpa" + + +@dataclass(frozen=True) +class _ResolvedAttention: + attention: NAttentionKind = NAttentionKind.EAGER_SDPA + compile_blocks: bool = False + + +def resolve_attention_for_host(mode): + del mode + return _ResolvedAttention() + + +def frames_per_yuv_gemm(height: int, width: int) -> int: + del height, width + return 2**31 - 1 + + +def patchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: + 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: + 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): + def __init__(self, latent_channels: int = 128): + super().__init__() + # Frozen parameters, not buffers: disk-offload reload restores only named_parameters(), + # so persistent buffers would be skipped and strict load_state_dict would report them missing. + self.register_parameter("std-of-means", nn.Parameter(torch.ones(latent_channels), requires_grad=False)) + self.register_parameter("mean-of-means", nn.Parameter(torch.zeros(latent_channels), requires_grad=False)) + + def un_normalize(self, x: torch.Tensor) -> torch.Tensor: + return (x * getattr(self, "std-of-means").view(1, -1, 1, 1, 1).to(x)) + getattr(self, "mean-of-means").view( + 1, -1, 1, 1, 1 + ).to(x) + + def normalize(self, x: torch.Tensor) -> torch.Tensor: + return (x - getattr(self, "mean-of-means").view(1, -1, 1, 1, 1).to(x)) / getattr(self, "std-of-means").view( + 1, -1, 1, 1, 1 + ).to(x) + + +"""Keyframe (dual-stream) inputs and coordinate math for DiffVAE decode. +A keyframe-aware decode carries two streams through the decoder: the video volume +``(B, T, H, W, C)`` and a stack of keyframe *planes* ``(B, P, H, W, C)`` whose plane +axis occupies video's temporal slot. Weights are fully shared; the streams only ever +mix inside one joint attention softmax (see ``transformer/fallback_na/joint_eager.py``). +Everything here is pure coordinate/geometry math with no module state, so the eager and +triton backends can share it and therefore agree exactly on slot selection. +Deviation from upstream worth knowing: upstream carries per-sample keyframe times and +masks (``(B, n_kf)``). Here they are batch-shared 1-D ``(P,)`` tensors, because our +decode path is single-sample and ``rope_math.rot_abs_axis_impl`` takes a 1-D position +vector per axis. That keeps the RoPE call and the slot tables batch-independent. +""" + + +KEYFRAME_CONTEXT_SLOTS = 2 + + +@dataclass(frozen=True) +class DecodeKeyframes: + latents: torch.Tensor + pixel_frame_indices: torch.Tensor + clip_start_frame: int = 0 + + def validate(self, *, num_frames: int | None = None) -> None: + if self.latents.ndim != 5: + raise ValueError(f"keyframe latents must be (B, C, P, H, W), got {tuple(self.latents.shape)}") + if self.pixel_frame_indices.ndim != 1: + raise ValueError(f"pixel_frame_indices must be 1-D (P,), got {tuple(self.pixel_frame_indices.shape)}") + if self.clip_start_frame < 0: + raise ValueError(f"clip_start_frame must be non-negative, got {self.clip_start_frame}") + planes = self.latents.shape[2] + if planes != self.pixel_frame_indices.shape[0]: + raise ValueError( + f"keyframe plane count {planes} != len(pixel_frame_indices) {self.pixel_frame_indices.shape[0]}" + ) + if planes == 0: + # An empty stack is a plain decode wearing a keyframe decode's costs, and every + # backend has to special-case it (the slot tables are all -1, and gathering plane 0 + # of an empty axis is an out-of-bounds read). Say so here instead. + raise ValueError("keyframe decode needs at least one plane; use decode_video() for a plain decode") + if planes and int(self.pixel_frame_indices.min()) < 0: + raise ValueError("pixel_frame_indices must be non-negative (global pixel frames)") + if num_frames is not None and num_frames < 1: + raise ValueError(f"num_frames must be positive, got {num_frames}") + # Planes may sit outside [clip_start_frame, clip_start_frame + num_frames): Dist tiles + # keep the nearest plane on each side so |dt| matches a whole-clip decode. A far plane + # on a full clip is the same geometry -- joint attention ranks it by distance. + + def crop_spatial(self, height: slice, width: slice) -> "DecodeKeyframes": + return DecodeKeyframes( + latents=self.latents[:, :, :, height, width], + pixel_frame_indices=self.pixel_frame_indices, + clip_start_frame=self.clip_start_frame, + ) + + @property + def num_planes(self) -> int: + return int(self.latents.shape[2]) + + +@dataclass(frozen=True) +class KeyframeStream: + x: torch.Tensor + times: torch.Tensor + valid: torch.Tensor + + def masked(self) -> KeyframeStream: + return KeyframeStream(x=self.x * self.valid[None, :, None, None, None], times=self.times, valid=self.valid) + + def select_planes(self, keep: torch.Tensor) -> KeyframeStream: + if keep.shape != (self.num_planes,): + raise ValueError(f"keep must be ({self.num_planes},) bool, got {tuple(keep.shape)}") + return KeyframeStream(x=self.x[:, keep], times=self.times[keep], valid=self.valid[keep]) + + def crop_spatial(self, height: slice, width: slice) -> KeyframeStream: + return KeyframeStream(x=self.x[:, :, height, width, :], times=self.times, valid=self.valid) + + @property + def num_planes(self) -> int: + return int(self.x.shape[1]) + + +def keyframe_stage_times(pixel_frame_indices: torch.Tensor, remaining_time_stride: int) -> torch.Tensor: + if remaining_time_stride < 1: + raise ValueError(f"remaining_time_stride must be positive, got {remaining_time_stride}") + frames = pixel_frame_indices.to(torch.float32) + center_offset = (remaining_time_stride - 1) / 2 + times = (frames + center_offset) / remaining_time_stride + return torch.where(frames == 0, torch.zeros_like(times), times) + + +def keyframe_clip_times( + pixel_frame_indices: torch.Tensor, + remaining_time_stride: int, + clip_start_frame: int, + extra_origin: float = 0.0, +) -> torch.Tensor: + times = keyframe_stage_times(pixel_frame_indices, remaining_time_stride) + origin = keyframe_stage_times( + torch.as_tensor([clip_start_frame], dtype=torch.int64, device=pixel_frame_indices.device), + remaining_time_stride, + ) + return times - origin - extra_origin + + +def planes_for_tile( + pixel_frame_indices: torch.Tensor, + frame_lo: int, + frame_hi: int, + *, + clip_start_frame: int = 0, +) -> torch.Tensor: + frame_lo = frame_lo + clip_start_frame + frame_hi = frame_hi + clip_start_frame + indices = pixel_frame_indices.to(torch.int64) + keep = (indices >= frame_lo) & (indices <= frame_hi) + before = indices < frame_lo + if bool(before.any()): + # Latest plane strictly before the tile. + keep[int(torch.where(before, indices, torch.full_like(indices, -1)).argmax())] = True + after = indices > frame_hi + if bool(after.any()): + # Earliest plane strictly after the tile. + sentinel = int(indices.max()) + 1 + keep[int(torch.where(after, indices, torch.full_like(indices, sentinel)).argmin())] = True + return keep + + +def remaining_time_strides(upsamples: Sequence[torch.nn.Module]) -> tuple[int, ...]: + strides = [int(up.stride[0]) for up in upsamples] + remaining: list[int] = [] + for index in range(len(strides)): + product = 1 + for stride in strides[index:]: + product *= stride + remaining.append(product) + remaining.append(1) + return tuple(remaining) + + +def upsample_keyframe_planes(upsample: torch.nn.Module, x: torch.Tensor) -> torch.Tensor: + planes = x.shape[1] + flat = rearrange(x, "b p h w c -> (b p) 1 h w c") + upsampled = upsample(flat, drop_leading_frame=True) + if upsampled.shape[1] != 1: + raise RuntimeError(f"isolated keyframe upsampling must preserve one temporal plane, got T={upsampled.shape[1]}") + out = rearrange(upsampled[:, 0], "(b p) h w c -> b p h w c", p=planes) + if out.shape[1] != planes: + raise RuntimeError(f"keyframe plane count changed under upsample: {planes} -> {out.shape[1]}") + return out + + +def _nearest_slots( + query_times: torch.Tensor, + candidate_times: torch.Tensor, + candidate_valid: torch.Tensor | None, + num_slots: int, +) -> torch.Tensor: + distances = (query_times[:, None] - candidate_times[None, :]).abs().to(torch.float32) + if candidate_valid is not None: + distances = distances.masked_fill(~candidate_valid[None, :], float("inf")) + order = torch.argsort(distances, dim=-1, stable=True) + take = min(num_slots, candidate_times.shape[0]) + chosen = order[:, :take] + # Drop slots that only exist because every remaining candidate was invalid. + finite = torch.gather(distances, 1, chosen).isfinite() + chosen = torch.where(finite, chosen, torch.full_like(chosen, -1)) + if take < num_slots: + pad = torch.full((chosen.shape[0], num_slots - take), -1, dtype=chosen.dtype, device=chosen.device) + chosen = torch.cat([chosen, pad], dim=1) + return chosen + + +def video_keyframe_slots( + keyframe_times: torch.Tensor, + keyframe_valid: torch.Tensor, + video_length: int, + num_slots: int = KEYFRAME_CONTEXT_SLOTS, +) -> torch.Tensor: + query = torch.arange(video_length, dtype=torch.float32, device=keyframe_times.device) + return _nearest_slots(query, keyframe_times.to(torch.float32), keyframe_valid, num_slots) + + +def keyframe_video_slots( + keyframe_times: torch.Tensor, + keyframe_valid: torch.Tensor, + video_length: int, + num_slots: int = KEYFRAME_CONTEXT_SLOTS, +) -> torch.Tensor: + candidates = torch.arange(video_length, dtype=torch.float32, device=keyframe_times.device) + slots = _nearest_slots(keyframe_times.to(torch.float32), candidates, None, num_slots) + return torch.where(keyframe_valid[:, None], slots, torch.full_like(slots, -1)) + + +"""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. +""" + + +ResizeAxisMode = Literal["repeat_last", "symmetric"] + + +_GIB: int = 1 << 30 + + +@dataclass(frozen=True, slots=True) +class _StageFiveBudget: + coef: float + coef_keyframes: float + reserve_bytes: int + + +_BUDGET_BY_MODE: dict[DiffVAEMode, _StageFiveBudget] = { + DiffVAEMode.COMBINED_COMPILE: _StageFiveBudget(coef=11, coef_keyframes=15, reserve_bytes=2 * _GIB), + DiffVAEMode.CHUNKED_COMPILE: _StageFiveBudget(coef=7, coef_keyframes=5, reserve_bytes=2 * _GIB), + DiffVAEMode.CHUNKED_EAGER: _StageFiveBudget(coef=5, coef_keyframes=5, reserve_bytes=1 * _GIB), + DiffVAEMode.BLACKWELL_DSL: _StageFiveBudget(coef=2.5, coef_keyframes=2.5, reserve_bytes=2 * _GIB), +} + + +_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 * _GIB + + +_BUDGET_SAFETY_BYTES_JOINT_MATERIALIZED: int = 2 * _GIB + + +def _falls_back_to_eager_na(mode: DiffVAEMode) -> bool: + resolved = resolve_attention_for_host(mode.resolve()) + return resolved.attention in (NAttentionKind.TRITON, NAttentionKind.EAGER_SDPA) and not resolved.compile_blocks + + +def stage5_mem_coef(mode: DiffVAEMode, *, keyframes: bool = False) -> float: + try: + budget = _BUDGET_BY_MODE[mode] + except KeyError as exc: + raise ValueError(f"Unsupported DiffVAEMode for tiling budget: {mode!r}") from exc + if keyframes: + return budget.coef_keyframes + if _falls_back_to_eager_na(mode): + return _BUDGET_BY_MODE[DiffVAEMode.CHUNKED_EAGER].coef + return budget.coef + + +_CONVERT_PEAK_UV_CHANNELS_X2 = 5 # uv_full (2) + pooled uv (0.5) + + +_PACK_PEAK_UV_CHANNELS_X2 = 4 # pooled uv (0.5) + packed float (1.5) + + +_PACK_PEAK_UINT8_BYTES_X2 = 3 # 1.5 B/px + + +def max_emitted_frames(*, num_frames: int, tile_frames: int, overlap_frames: int) -> int: + if tile_frames >= num_frames: + return num_frames + stride = tile_frames - overlap_frames + if stride <= 0: + return tile_frames + n_tiles = 1 + -(-(num_frames - tile_frames) // stride) + last_group = num_frames - (n_tiles - 1) * stride + return min(tile_frames, max(stride, last_group + 1)) + + +def emit_convert_bytes( + *, + tile_frames: int, + height: int, + width: int, + out_channels: int, + element_size: int, +) -> int: + frames = int(tile_frames) + gemm_frames = min(frames, frames_per_yuv_gemm(height, width)) + # Write-back reuses the RGB storage, so no full YUV tensor survives into the pack. + resident_yuv_frames = frames if gemm_frames == frames else 0 + yuv_channels_x2 = 2 * int(out_channels) + + convert_peak_x2 = element_size * (yuv_channels_x2 * gemm_frames + _CONVERT_PEAK_UV_CHANNELS_X2 * frames) + pack_peak_x2 = ( + element_size * (yuv_channels_x2 * resident_yuv_frames + _PACK_PEAK_UV_CHANNELS_X2 * frames) + + _PACK_PEAK_UINT8_BYTES_X2 * frames + ) + + return int(height) * int(width) * max(convert_peak_x2, pack_peak_x2) // 2 + + +def budget_safety_bytes( + mode: DiffVAEMode, + *, + keyframes: bool = False, + joint_sdpa_materializes: bool = False, +) -> int: + try: + budget = _BUDGET_BY_MODE[mode] + except KeyError as exc: + raise ValueError(f"Unsupported DiffVAEMode for tiling budget: {mode!r}") from exc + if keyframes and mode is not DiffVAEMode.BLACKWELL_DSL: + return _BUDGET_SAFETY_BYTES_JOINT_MATERIALIZED if joint_sdpa_materializes else _BUDGET_SAFETY_BYTES_EAGER + if _falls_back_to_eager_na(mode): + return _BUDGET_SAFETY_BYTES_EAGER + return budget.reserve_bytes + + +def accumulator_element_size(feature_dtype: torch.dtype) -> int: + 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: + 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 ltx_core.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) + + +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, + keyframes: bool = False, + joint_sdpa_materializes: bool = False, +) -> TileSizeConfig: + 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, keyframes=keyframes) + 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, + ) + reserve = budget_safety_bytes(mode, keyframes=keyframes, joint_sdpa_materializes=joint_sdpa_materializes) + usable = max(0, int(free_bytes) - model_cost - reserve - 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) + + 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 + # The consumer converts each yielded chunk while this decode is suspended. + downstream_bytes = emit_convert_bytes( + tile_frames=max_emitted_frames(num_frames=num_frames, tile_frames=tile_t, overlap_frames=overlap_t), + height=height, + width=width, + out_channels=out_channels, + element_size=element_size, + ) + if acc_bytes + downstream_bytes >= usable: + continue + s5_budget_bytes = usable - acc_bytes - downstream_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}, keyframes={keyframes}, coef={coef}, stage5_channels={stage5_channels}, " + f"stage4_feature_bytes={s4_feat_bytes}, usable_bytes={usable}. " + "Reduce resolution, reduce num_frames (stage-4 features and the per-chunk " + "encode buffers both scale with it), 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]: + 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]]: + 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 + + +@dataclass(frozen=True) +class AxisPad: + before: int + after: int + + +def resize_axis( + x: torch.Tensor, + dim: int, + size: int, + *, + mode: ResizeAxisMode, +) -> tuple[torch.Tensor, AxisPad]: + 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]]: + 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: + 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: + 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]: + 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: + 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: + 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: + 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]: + 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: + 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]: + 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]]: + 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]]: + 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]: + 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) + + +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]: + 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: + 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]]: + 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 = n_t * n_h * n_w * tile_frames * tile_height * tile_width + unique = max(1, num_frames * height * width) + return processed / unique + + +def _propagate_interval_through_upsample_hops( + interval: DimensionInterval, + strides: Sequence[int], + causal: bool, +) -> DimensionInterval: + 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 + + +"""Shared small layers for the diffusion-VAE NA transformer stack.""" + + +class ChannelLinear(nn.Linear): + @property + def in_channels(self) -> int: + return self.in_features + + @property + def out_channels(self) -> int: + return self.out_features + + +class LinearPixelShuffleUpsample(nn.Module): + 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: + 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): + 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: + return x * (1.0 + scale) + shift + + +"""Shared absolute-RoPE math helpers (no consumer policy).""" + + +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]: + 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: + 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: + 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 + + +"""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. +""" + + +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, + t_pos: torch.Tensor | None = None, +) -> torch.Tensor: + d_t, d_h, _ = rope_split + inv_t, inv_h, inv_w = inv_freqs + t = x.shape[1] + h = x.shape[2] + positions_t = t_positions(t, x.device) if t_pos is None else t_pos + xt = rot_abs_axis_impl(x[..., :d_t], positions_t, 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, + t_pos: torch.Tensor | None = None, +) -> torch.Tensor: + 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, + t_pos=t_pos, + ) + ) + w_off = w_off + w_slab + return torch.cat(parts, dim=3) + + +@torch.library.custom_op("diffsynth_ltx25::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: + 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, + ) + + +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, + ) + + +@torch.library.custom_op("diffsynth_ltx25::abs_rope_at_t", mutates_args=()) +def _abs_rope_at_t_op( + x: torch.Tensor, + t_pos: 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: + 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, + t_pos=t_pos, + ) + + +def _rope_config(attn: object, x: torch.Tensor) -> tuple[tuple[torch.Tensor, ...], int, torch.dtype]: + 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) + return inv_freqs, num_tiles, compute_dtype + + +def _det_project_qkv(attn: object, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + 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] + return q, k, v + + +def det_qkv_rope_at_times( + attn: object, + x: torch.Tensor, + t_pos: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if t_pos.ndim != 1 or t_pos.shape[0] != x.shape[1]: + raise ValueError(f"t_pos must be ({x.shape[1]},) to match the plane axis, got {tuple(t_pos.shape)}") + q, k, v = _det_project_qkv(attn, x) + + inv_freqs, num_tiles, compute_dtype = _rope_config(attn, x) + d_t, d_h, d_w = attn.rope_dim_split # type: ignore[attr-defined] + positions = t_pos.to(device=x.device, dtype=torch.float32) + rotated = [ + _abs_rope_at_t_op( + tensor, + positions, + *inv_freqs, + d_t, + d_h, + d_w, + num_tiles, + compute_dtype == torch.bfloat16, + ) + for tensor in (q, k) + ] + return rotated[0], rotated[1], v + + +def det_qkv_rope(attn: object, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, v = _det_project_qkv(attn, x) + inv_freqs, num_tiles, compute_dtype = _rope_config(attn, x) + 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 + + +def vram_ready_linear(module: nn.Module) -> tuple[torch.Tensor, torch.Tensor | None]: + # This decoder calls several projections functionally, bypassing the VRAM wrappers' + # forward, so ask the wrapper for computation-ready weights instead of reading them raw. + computation = getattr(module, "computation", None) + if computation is not None: + return computation() + return module.weight, module.bias + + +class QKVProjections(nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + linear = nn.Linear(dim, dim * 3, bias=True) + self.weight = linear.weight + self.bias = linear.bias + + def forward(self, x): + weight = self.weight.to(device=x.device, dtype=x.dtype) + bias = self.bias.to(device=x.device, dtype=x.dtype) + weights = weight.chunk(3, dim=0) + biases = bias.chunk(3, dim=0) + return tuple(F.linear(x, weight, bias) for weight, bias in zip(weights, biases, strict=True)) + + +DEFAULT_SWIGLU_TILE_SIZE: Final[int] = 16_384 +DEFAULT_SWIGLU_TILES: Final[int] = 4 + + +@dataclass(frozen=True) +class SwiGLUTileSpec: + 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("num_tiles must be >= 1") + if self.tile_size is not None and self.tile_size < 1: + raise ValueError("tile_size must be >= 1") + +DEFAULT_SWIGLU_TILE_SPEC = SwiGLUTileSpec(tile_size=DEFAULT_SWIGLU_TILE_SIZE) + + +def _swiglu_chunk(x, w_gate, w_up, w_down): + return F.linear(F.silu(F.linear(x, w_gate)) * F.linear(x, w_up), w_down) + + +def swiglu_tiled(x, w_gate, w_up, w_down, tile, *, use_triton=None): + del use_triton + if x.numel() == 0: + return x + leading = x.shape[:-1] + flat = x.reshape(-1, x.shape[-1]) + if tile.tile_size is not None: + chunk_size = tile.tile_size + else: + chunk_size = max(1, math.ceil(flat.shape[0] / tile.num_tiles)) + output = torch.cat([_swiglu_chunk(chunk, w_gate, w_up, w_down) for chunk in flat.split(chunk_size)], dim=0) + return output.reshape(*leading, output.shape[-1]) + + +def swiglu_weights(mlp) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return ( + vram_ready_linear(mlp.w_gate)[0], + vram_ready_linear(mlp.w_up)[0], + vram_ready_linear(mlp.w_down)[0], + ) + + +def plain_mlp(x, mlp, norm, tile): + y = norm(x) + if y.numel() == 0: + return x + return x + swiglu_tiled(y, *swiglu_weights(mlp), tile) + + +class SwiGLU(nn.Module): + 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): + return swiglu_tiled(x, *swiglu_weights(self), self.tile) + + +"""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. +""" + + +NA_SCORE_BUDGET = 2**25 + + +NA_KV_STACK_BUDGET = 2**28 + + +def _window_bounds(length: int, kernel: int, causal: bool) -> tuple[list[int], list[int]]: + 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]: + 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: + 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: + 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 + + +"""Pure-torch joint (video + keyframe) 3D neighborhood attention. +Computes, for one softmax per query: +* **video query** at ``(t, h, w)``: its local ``Kt x Kh x Kw`` video window, clamped to the + volume and masked where it hangs over the edge, **plus** the whole ``Kh x Kw`` window at the + same ``(h, w)`` on each of the ``num_slots`` nearest keyframe planes. Keyframe visibility does + not depend on ``Kt`` -- a plane far outside the temporal radius is still visible. +* **keyframe query** on plane ``i``: the ``Kh x Kw`` window on its own plane (there is no + plane-to-plane attention), plus the same window on each of the nearest video frames. +Which planes and frames are "nearest" comes from :func:`video_keyframe_slots` / +:func:`keyframe_video_slots`, so every backend agrees on visibility. +No Triton, no natten, no ``torch.compile``: this is the backend that always exists -- CPU, +macOS/MPS, Windows without a built extra. +Structure +--------- +Everything is arranged so the arithmetic happens inside ``F.scaled_dot_product_attention``: +**Query bricks.** Queries are grouped into ``(bt, bh, bw)`` bricks and many bricks ride one SDPA +call as its batch dimension, so a frame costs a handful of launches rather than thousands. All +queries in a brick share one gathered key slab, of extent ``(bt + Kt - 1, bh + Kh - 1, bw + Kw - 1)``. +**Brick shape.** Wasted work is ``Nk / keys_actually_visible`` and *grows* with the brick, so the +spatial face stays small and square (square minimizes the slab at a fixed query count). Depth is +the exception: the gather is the larger cost and it scales with ``Nk / Nq``, which *falls* with +depth, so :data:`DEFAULT_BRICK_DEPTH` frames deep beats one frame deep despite doing more +arithmetic. Both defaults sit on measured plateaus at the production stage-5 shape. +**One shared, 2D-broadcast mask.** The visible-key pattern is a property of the brick geometry, +identical for every brick, so it is built once and passed as ``(1, 1, Nq, Nk)``. That shape is +load-bearing: torch keeps the memory-efficient backend and expands neither the mask nor the +scores, whereas a pre-expanded ``(G, NH, Nq, Nk)`` bias halves throughput and costs gigabytes. +**Per-key validity rides in the keys.** Out-of-volume positions, empty slots (``-1``) and invalid +planes are data-dependent, so folding them into the mask would make it per-brick. Instead ``K`` +carries one extra channel holding ``0`` for a live key and :data:`_DEAD` for a dead one, against a +constant ``1`` channel on ``Q``. Q arrives pre-scaled, so with ``scale=1.0`` that adds exactly the +bias to the score. The channel count is then rounded up to :data:`_HEAD_DIM_ALIGN`. +**Head-major staging.** Key slabs are gathered from a ``(B, NH, A, Hp, Wp, C)`` copy rather than +from the caller's channels-last layout, so the gather's innermost contiguous run is ``ew * C`` +instead of ``C``. +**Runs of constant slot row.** A brick spanning several frames shares one keyframe key slab, so it +must not straddle a change of visible planes. ``T`` is cut into maximal runs of identical slot rows +and bricks are tiled inside a run -- which also means one plane gather per run, entering the slab +view with a **zero** group stride. +Both loops are budgeted by :data:`DEFAULT_WORKSPACE_BYTES`: frames per staging pass, then +``(bricks, brick rows)`` per SDPA call. Peak transient memory is therefore bounded by that budget +and not by the volume, which is what lets this sit next to a decoder that has its own memory plan. +The gather is the floor: SDPA needs materialized ``(G, NH, Nk, HD)`` keys, so every key is copied +``Nk / Nq`` times. Only a fused neighborhood kernel avoids that. +""" + + +_DEAD = -1.0e4 + + +_HEAD_DIM_ALIGN = 8 + + +DEFAULT_BRICK_QUERIES = 64 + + +DEFAULT_BRICK_DEPTH = 4 + + +DEFAULT_WORKSPACE_BYTES = 256 * 1024**2 + + +_STAGING_FACTOR_FUSED = 4.75 + + +_STAGING_FACTOR_MATERIALIZED = 22.1 + + +def sdpa_materializes_scores(device: torch.device) -> bool: + return device.type != "cuda" + + +def staging_factor(device: torch.device) -> float: + return _STAGING_FACTOR_MATERIALIZED if sdpa_materializes_scores(device) else _STAGING_FACTOR_FUSED + + +def _key_channels(head_dim: int) -> int: + return -(-(head_dim + 1) // _HEAD_DIM_ALIGN) * _HEAD_DIM_ALIGN + + +def _window(kernel: int) -> tuple[int, int]: + lo = kernel // 2 + return lo, kernel - lo - 1 + + +def pick_brick( + time: int, + height: int, + width: int, + target: int = DEFAULT_BRICK_QUERIES, + depth: int = DEFAULT_BRICK_DEPTH, +) -> tuple[int, int, int]: + side = max(1, round(math.sqrt(target))) + return min(depth, time), min(side, height), min(side, width) + + +class _Geometry: + def __init__( + self, + height: int, + width: int, + kernel: tuple[int, int, int], + brick: tuple[int, int, int], + ) -> None: + kernel_t, kernel_h, kernel_w = kernel + lo_h, hi_h = _window(kernel_h) + lo_w, hi_w = _window(kernel_w) + self.height, self.width = height, width + self.brick = brick + self.kernel = kernel + self.grid = (-(-height // brick[1]), -(-width // brick[2])) + # Slab extents: T grows with the brick depth, H/W with the spatial face. + self.span_t = brick[0] + kernel_t - 1 + self.span = (brick[1] + kernel_h - 1, brick[2] + kernel_w - 1) + # Halo, plus enough to cover the last (partial) brick's slab. + self.pad_h = (lo_h, hi_h + self.grid[0] * brick[1] - height) + self.pad_w = (lo_w, hi_w + self.grid[1] * brick[2] - width) + self.pad_t = _window(kernel_t) + self.queries = brick[0] * brick[1] * brick[2] + self.footprint = self.span[0] * self.span[1] + self.padded_height = height + sum(self.pad_h) + self.padded_width = width + sum(self.pad_w) + + def row_extent(self, rows: int) -> int: + return (rows - 1) * self.brick[1] + self.span[0] + + +class _Schedule: + def __init__( + self, + geometry: _Geometry, + blocks: int, + heads: int, + head_dim: int, + axis_bricks: int, + element_size: int, + workspace_bytes: int, + factor: float, + ) -> None: + channels = _key_channels(head_dim) + # One (brick along the axis, brick row) pair's worth of gathered keys and values, plus its + # score block on backends that materialize one. + keys = blocks * geometry.footprint + pair_bytes = geometry.grid[1] * heads * keys * (channels + head_dim) * element_size + # ``factor`` folds in whatever the selected SDPA kernel allocates on top of the staging, + # chiefly a materialized score block. See :data:`_STAGING_FACTOR_FUSED`. + pairs = max(1, int(workspace_bytes / max(pair_bytes * factor, 1.0))) + if pairs >= geometry.grid[0]: + self.group_axis = min(axis_bricks, max(1, pairs // geometry.grid[0])) + self.group_rows = geometry.grid[0] + else: + self.group_axis = 1 + self.group_rows = pairs + staged = geometry.padded_height * geometry.padded_width * heads * (channels + head_dim) * element_size + per_axis_brick = staged * geometry.brick[0] + self.stage_axis = min(axis_bricks, max(self.group_axis, workspace_bytes // max(per_axis_brick, 1))) + + +def _banded(queries: int, span: int, kernel: int, device: torch.device) -> torch.Tensor: + key = torch.arange(span, device=device)[None, :] + query = torch.arange(queries, device=device)[:, None] + return (key >= query) & (key < query + kernel) + + +def _joint_mask(geometry: _Geometry, num_slots: int, device: torch.device) -> torch.Tensor: + brick_t, brick_h, brick_w = geometry.brick + kernel_t, kernel_h, kernel_w = geometry.kernel + spatial = ( + _banded(brick_h, geometry.span[0], kernel_h, device)[:, None, :, None] + & _banded(brick_w, geometry.span[1], kernel_w, device)[None, :, None, :] + ).reshape(brick_h * brick_w, geometry.footprint) + temporal = _banded(brick_t, geometry.span_t, kernel_t, device) + video = (temporal[:, None, :, None] & spatial[None, :, None, :]).reshape( + geometry.queries, geometry.span_t * geometry.footprint + ) + planes = ( + spatial[None, :, None, :] + .expand(brick_t, brick_h * brick_w, num_slots, geometry.footprint) + .reshape(geometry.queries, num_slots * geometry.footprint) + ) + return torch.cat([video, planes], dim=1)[None, None].contiguous() + + +def _stage( + x: torch.Tensor, + geometry: _Geometry, + pad_t: tuple[int, int], + *, + with_bias_channel: bool, +) -> torch.Tensor: + batch, axis, height, width, heads, head_dim = x.shape + channels = _key_channels(head_dim) if with_bias_channel else head_dim + out = x.new_zeros((batch, heads, axis + sum(pad_t), geometry.padded_height, geometry.padded_width, channels)) + if with_bias_channel: + out[..., head_dim] = _DEAD + live = out[ + :, + :, + pad_t[0] : pad_t[0] + axis, + geometry.pad_h[0] : geometry.pad_h[0] + height, + geometry.pad_w[0] : geometry.pad_w[0] + width, + ] + live[..., :head_dim] = x.permute(0, 4, 1, 2, 3, 5) + if with_bias_channel: + live[..., head_dim] = 0.0 + return out + + +def _slabs( + staged: torch.Tensor, + geometry: _Geometry, + bricks: int, + rows: int, + blocks: int, + *, + group_stride: int, +) -> torch.Tensor: + batch, heads = staged.shape[0], staged.shape[1] + stride_b, stride_nh, stride_a, stride_h, stride_w, _ = staged.stride() + return staged.as_strided( + (batch, bricks, rows, geometry.grid[1], heads, blocks, *geometry.span, staged.shape[-1]), + ( + stride_b, + group_stride * stride_a, + geometry.brick[1] * stride_h, + geometry.brick[2] * stride_w, + stride_nh, + stride_a, + stride_h, + stride_w, + 1, + ), + ) + + +def _query_bricks(x: torch.Tensor, geometry: _Geometry, bricks: int, rows: int) -> torch.Tensor: + batch, axis, height, width, heads, head_dim = x.shape + brick_t, brick_h, brick_w = geometry.brick + pad_t, pad_h, pad_w = bricks * brick_t - axis, rows * brick_h - height, geometry.grid[1] * brick_w - width + if pad_t or pad_h or pad_w: + x = F.pad(x, (0, 0, 0, 0, 0, pad_w, 0, pad_h, 0, pad_t)) + bricked = ( + x.reshape(batch, bricks, brick_t, rows, brick_h, geometry.grid[1], brick_w, heads, head_dim) + .permute(0, 1, 3, 5, 7, 2, 4, 6, 8) + .reshape(batch * bricks * rows * geometry.grid[1], heads, geometry.queries, head_dim) + ) + out = bricked.new_zeros((*bricked.shape[:-1], _key_channels(head_dim))) + out[..., :head_dim] = bricked + out[..., head_dim] = 1.0 + return out + + +def _unbrick( + attended: torch.Tensor, + geometry: _Geometry, + batch: int, + bricks: int, + rows: int, + extent: tuple[int, int], +) -> torch.Tensor: + brick_t, brick_h, brick_w = geometry.brick + heads, head_dim = attended.shape[1], attended.shape[3] + plane = ( + attended.reshape(batch, bricks, rows, geometry.grid[1], heads, brick_t, brick_h, brick_w, head_dim) + .permute(0, 1, 5, 2, 6, 3, 7, 4, 8) + .reshape(batch, bricks * brick_t, rows * brick_h, geometry.grid[1] * brick_w, heads, head_dim) + ) + return plane[:, : extent[0], : extent[1], : geometry.width] + + +def _with_null(slots: torch.Tensor, null_index: int) -> torch.Tensor: + return torch.where(slots < 0, torch.full_like(slots, null_index), slots) + + +def _append_null(keys: torch.Tensor, values: torch.Tensor, head_dim: int) -> tuple[torch.Tensor, torch.Tensor]: + shape = (keys.shape[0], keys.shape[1], 1, *keys.shape[3:]) + null_key = keys.new_zeros(shape) + null_key[..., head_dim] = _DEAD + null_value = values.new_zeros((*shape[:-1], values.shape[-1])) + return torch.cat([keys, null_key], dim=2), torch.cat([values, null_value], dim=2) + + +def _slot_runs(slots: torch.Tensor) -> list[tuple[int, int]]: + rows = slots.tolist() + runs: list[tuple[int, int]] = [] + start = 0 + for index in range(1, len(rows)): + if rows[index] != rows[start]: + runs.append((start, index)) + start = index + runs.append((start, len(rows))) + return runs + + +def _attend_group( + query_slice: torch.Tensor, + key_views: tuple[torch.Tensor, ...], + value_views: tuple[torch.Tensor, ...], + geometry: _Geometry, + shape: tuple[int, int], + mask: torch.Tensor, +) -> torch.Tensor: + bricks, rows = shape + batch = query_slice.shape[0] + heads, head_dim = query_slice.shape[4], query_slice.shape[5] + blocks = sum(view.shape[5] for view in key_views) + channels = _key_channels(head_dim) + keys = query_slice.new_empty((batch, bricks, rows, geometry.grid[1], heads, blocks, *geometry.span, channels)) + values = query_slice.new_empty((batch, bricks, rows, geometry.grid[1], heads, blocks, *geometry.span, head_dim)) + start = 0 + for key_view, value_view in zip(key_views, value_views, strict=True): + stop = start + key_view.shape[5] + keys[:, :, :, :, :, start:stop].copy_(key_view) + values[:, :, :, :, :, start:stop].copy_(value_view) + start = stop + count = batch * bricks * rows * geometry.grid[1] + attended = F.scaled_dot_product_attention( + _query_bricks(query_slice, geometry, bricks, rows), + keys.view(count, heads, blocks * geometry.footprint, channels), + values.view(count, heads, blocks * geometry.footprint, head_dim), + attn_mask=mask, + scale=1.0, + ) + return _unbrick(attended, geometry, batch, bricks, rows, (query_slice.shape[1], query_slice.shape[2])) + + +def _row_groups(geometry: _Geometry, schedule: _Schedule) -> list[tuple[int, int, slice, slice]]: + brick_h = geometry.brick[1] + groups = [] + for row in range(0, geometry.grid[0], schedule.group_rows): + rows = min(schedule.group_rows, geometry.grid[0] - row) + groups.append( + ( + row, + rows, + slice(row * brick_h, row * brick_h + geometry.row_extent(rows)), + slice(row * brick_h, min((row + rows) * brick_h, geometry.height)), + ) + ) + return groups + + +def _video_query_pass( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + keyframe_k: torch.Tensor, + keyframe_v: torch.Tensor, + slots: torch.Tensor, + geometry: _Geometry, + workspace_bytes: int, + factor: float, +) -> torch.Tensor: + time, heads, head_dim = q.shape[1], q.shape[4], q.shape[5] + brick_t = geometry.brick[0] + lo_t, hi_t = geometry.pad_t + num_slots = slots.shape[1] + blocks = geometry.span_t + num_slots + + plane_keys, plane_values = _append_null( + _stage(keyframe_k, geometry, (0, 0), with_bias_channel=True), + _stage(keyframe_v, geometry, (0, 0), with_bias_channel=False), + head_dim, + ) + slot_table = _with_null(slots, keyframe_k.shape[1]) + mask = _joint_mask(geometry, num_slots, q.device) + schedule = _Schedule( + geometry, + blocks, + heads, + head_dim, + -(-time // brick_t), + q.element_size(), + workspace_bytes, + factor, + ) + rows_groups = _row_groups(geometry, schedule) + + out = torch.empty_like(q) + for run_start, run_stop in _slot_runs(slot_table): + # One plane gather per run: every brick inside it sees the same slots. + planes = plane_keys.index_select(2, slot_table[run_start]) + plane_vals = plane_values.index_select(2, slot_table[run_start]) + run_bricks = -(-(run_stop - run_start) // brick_t) + for staged_brick in range(0, run_bricks, schedule.stage_axis): + staged_bricks = min(schedule.stage_axis, run_bricks - staged_brick) + first = run_start + staged_brick * brick_t + last = first + staged_bricks * brick_t # exclusive; may reach past the run or T + source = slice(max(0, first - lo_t), min(time, last + hi_t)) + pad_t = (max(0, lo_t - first), max(0, last + hi_t - time)) + window_keys = _stage(k[:, source], geometry, pad_t, with_bias_channel=True) + window_values = _stage(v[:, source], geometry, pad_t, with_bias_channel=False) + + for brick in range(staged_brick, staged_brick + staged_bricks, schedule.group_axis): + count = min(schedule.group_axis, staged_brick + staged_bricks - brick) + start = run_start + brick * brick_t + stop = min(start + count * brick_t, run_stop) + offset = (brick - staged_brick) * brick_t + for _, rows, key_rows, out_rows in rows_groups: + tile = _attend_group( + q[:, start:stop, out_rows], + ( + _slabs( + window_keys[:, :, offset:, key_rows], + geometry, + count, + rows, + geometry.span_t, + group_stride=brick_t, + ), + _slabs(planes[:, :, :, key_rows], geometry, count, rows, num_slots, group_stride=0), + ), + ( + _slabs( + window_values[:, :, offset:, key_rows], + geometry, + count, + rows, + geometry.span_t, + group_stride=brick_t, + ), + _slabs(plane_vals[:, :, :, key_rows], geometry, count, rows, num_slots, group_stride=0), + ), + geometry, + (count, rows), + mask, + ) + out[:, start:stop, out_rows] = tile + return out + + +def _keyframe_query_pass( + keyframe_q: torch.Tensor, + keyframe_k: torch.Tensor, + keyframe_v: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + slots: torch.Tensor, + keyframe_valid: torch.Tensor, + geometry: _Geometry, + workspace_bytes: int, + factor: float, +) -> torch.Tensor: + planes_total, heads, head_dim = keyframe_q.shape[1], keyframe_q.shape[4], keyframe_q.shape[5] + num_slots = slots.shape[1] + blocks = 1 + num_slots + time = k.shape[1] + flat = _Geometry(geometry.height, geometry.width, (1, *geometry.kernel[1:]), (1, *geometry.brick[1:])) + + # Only the frames some plane actually points at get staged -- at most ``P * num_slots`` of them, + # against the whole volume if this staged ``k`` wholesale. ``unique`` doubles as the remap: slot + # rows are rewritten to index the compacted stack. + wanted, inverse = torch.unique(_with_null(slots, time).reshape(-1), return_inverse=True) + frame_keys = _stage(k.index_select(1, wanted.clamp(max=time - 1)), flat, (0, 0), with_bias_channel=True) + frame_values = _stage(v.index_select(1, wanted.clamp(max=time - 1)), flat, (0, 0), with_bias_channel=False) + # An empty slot clamped onto a real frame above; kill it here instead of appending a null row. + frame_keys[:, :, wanted == time, ..., head_dim] = _DEAD + own_keys = _stage(keyframe_k, flat, (0, 0), with_bias_channel=True) + own_values = _stage(keyframe_v, flat, (0, 0), with_bias_channel=False) + own_keys[:, :, ~keyframe_valid, ..., head_dim] = _DEAD + slot_table = inverse.reshape(planes_total, num_slots) + mask = _joint_mask(flat, num_slots, keyframe_q.device) + schedule = _Schedule( + flat, + blocks, + heads, + head_dim, + planes_total, + keyframe_q.element_size(), + workspace_bytes, + factor, + ) + rows_groups = _row_groups(flat, schedule) + + out = torch.empty_like(keyframe_q) + for start in range(0, planes_total, schedule.group_axis): + stop = min(start + schedule.group_axis, planes_total) + count = stop - start + picked = slot_table[start:stop].reshape(-1) + frames = frame_keys.index_select(2, picked) + frame_vals = frame_values.index_select(2, picked) + for _, rows, key_rows, out_rows in rows_groups: + tile = _attend_group( + keyframe_q[:, start:stop, out_rows], + ( + _slabs(own_keys[:, :, start:, key_rows], flat, count, rows, 1, group_stride=1), + _slabs(frames[:, :, :, key_rows], flat, count, rows, num_slots, group_stride=num_slots), + ), + ( + _slabs(own_values[:, :, start:, key_rows], flat, count, rows, 1, group_stride=1), + _slabs(frame_vals[:, :, :, key_rows], flat, count, rows, num_slots, group_stride=num_slots), + ), + flat, + (count, rows), + mask, + ) + out[:, start:stop, out_rows] = tile + # An invalid plane sees nothing; zero it rather than shipping the uniform mean. + return out * keyframe_valid[None, :, None, None, None, None] + + +def joint_na3d( # noqa: PLR0913 + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + keyframe_q: torch.Tensor, + keyframe_k: torch.Tensor, + keyframe_v: torch.Tensor, + keyframe_times: torch.Tensor, + keyframe_valid: torch.Tensor, + kernel_size: tuple[int, int, int], + num_slots: int = KEYFRAME_CONTEXT_SLOTS, + brick: tuple[int, int, int] | None = None, + workspace_bytes: int = DEFAULT_WORKSPACE_BYTES, + factor: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + time, height, width = q.shape[1], q.shape[2], q.shape[3] + video_slots = video_keyframe_slots(keyframe_times, keyframe_valid, time, num_slots) + keyframe_slots = keyframe_video_slots(keyframe_times, keyframe_valid, time, num_slots) + geometry = _Geometry(height, width, kernel_size, brick if brick is not None else pick_brick(time, height, width)) + if factor is None: + factor = staging_factor(q.device) + return ( + _video_query_pass(q, k, v, keyframe_k, keyframe_v, video_slots, geometry, workspace_bytes, factor), + _keyframe_query_pass( + keyframe_q, + keyframe_k, + keyframe_v, + k, + v, + keyframe_slots, + keyframe_valid, + geometry, + workspace_bytes, + factor, + ), + ) + + +class EagerNAAttention: + def __call__(self, attn, q, k, v): + return na3d(q, k, v, kernel_size=attn.kernel_size, scale=1.0) + + +class EagerJointNAAttention: + def __call__(self, attn, q, k, v, keyframe_q, keyframe_k, keyframe_v, keyframe_times, keyframe_valid): + return joint_na3d( + q, k, v, keyframe_q, keyframe_k, keyframe_v, + keyframe_times, keyframe_valid, kernel_size=attn.kernel_size, + ) + + +"""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). +""" + + +if TYPE_CHECKING: + from ltx_core.model.video_vae.keyframes import KeyframeStream + + +try: + import natten + + _NATTEN_AVAILABLE = True +except ImportError: # pragma: no cover + natten = None # type: ignore[assignment] + _NATTEN_AVAILABLE = False + + +class NAAttentionCallable(Protocol): + def __call__( + self, + attn: NeighborhoodAttention3D, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> torch.Tensor: ... + + +class JointNAAttentionCallable(Protocol): + def __call__( + self, + attn: NeighborhoodAttention3D, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + keyframe_q: torch.Tensor, + keyframe_k: torch.Tensor, + keyframe_v: torch.Tensor, + keyframe_times: torch.Tensor, + keyframe_valid: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: ... + + +class NeighborhoodAttention3D(nn.Module): + 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 = EagerNAAttention() + # Separate slot, installed for every mode; never NATTEN/DSL. Only the keyframe + # decode path reads it, so keyframe-less decode keeps NATTEN when it is installed. + self.joint_attention_function: JointNAAttentionCallable | None = EagerJointNAAttention() + + 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]: + 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: + 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 forward_with_keyframes( + self, + x: torch.Tensor, + keyframes: KeyframeStream, + ) -> tuple[torch.Tensor, KeyframeStream]: + if self.joint_attention_function is None: + raise RuntimeError( + "keyframe decode needs joint_attention_function installed; build the decoder " + "through apply_diffvae_config / apply_diffvae_mode" + ) + batch, t, h, w, _ = x.shape + planes = keyframes.x.shape[1] + + q, k, v = det_qkv_rope(self, x) + keyframe_q, keyframe_k, keyframe_v = det_qkv_rope_at_times(self, keyframes.x, keyframes.times) + q, k, v = q.contiguous(), k.contiguous(), v.contiguous() + keyframe_q = keyframe_q.contiguous() + keyframe_k = keyframe_k.contiguous() + keyframe_v = keyframe_v.contiguous() + + out, keyframe_out = self.joint_attention_function( + self, + q, + k, + v, + keyframe_q, + keyframe_k, + keyframe_v, + keyframes.times, + keyframes.valid, + ) + out = self.proj(out.reshape(batch, t, h, w, self.dim)) + keyframe_out = self.proj(keyframe_out.reshape(batch, planes, h, w, self.dim)) + return out, dataclasses.replace(keyframes, x=keyframe_out) + + +"""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. +""" + + +if TYPE_CHECKING: + from ltx_core.model.video_vae.keyframes import KeyframeStream + + +__all__ = [ + "DiffusionNABlock", + "NABlock", +] + + +class NABlock(nn.Module): + 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: + x = x + self.attn(self.norm1(x)) + x = plain_mlp(x, self.mlp, self.norm2, self.mlp.tile) + return x + + def forward_with_keyframes( + self, + x: torch.Tensor, + keyframes: KeyframeStream, + ) -> tuple[torch.Tensor, KeyframeStream]: + attn_out, keyframe_attn = self.attn.forward_with_keyframes( + self.norm1(x), + dataclasses.replace(keyframes, x=self.norm1(keyframes.x)), + ) + x = x + attn_out + keyframe_x = keyframes.x + keyframe_attn.x + x = plain_mlp(x, self.mlp, self.norm2, self.mlp.tile) + keyframe_x = plain_mlp(keyframe_x, self.mlp, self.norm2, self.mlp.tile) + return x, dataclasses.replace(keyframes, x=keyframe_x) + + +class DiffusionNABlock(nn.Module): + 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_shift_table = self.scale_shift_table.to(dtype=modulation[0].dtype, device=modulation[0].device) + scale_msa, shift_msa, _, scale_mlp, shift_mlp, _, _ = [ + modulation[i] + 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 + + +"""Combined context residual: project context half of ``context_and_x`` into ``x``.""" + + +def combined( + context_and_x: torch.Tensor, + w_proj: torch.Tensor, + b_proj: torch.Tensor | None, +) -> torch.Tensor: + 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) + + +inject_context = combined + + +"""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. +""" + + +_rot_abs_axis = torch.compiler.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, + t_pos: torch.Tensor | None = None, +) -> torch.Tensor: + d_t, d_h, _ = rope_split + inv_t, inv_h, inv_w = inv_freqs + t = x.shape[1] + h = x.shape[2] + positions_t = t_positions(t, x.device) if t_pos is None else t_pos + xt = _rot_abs_axis(x[..., :d_t], positions_t, 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, + t_pos: torch.Tensor | None = None, +) -> torch.Tensor: + 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, + t_pos=t_pos, + ) + ) + w_off = w_off + w_slab + return torch.cat(parts, dim=3) + + +def _qkv_nested_rope( + attn: NeighborhoodAttention3D, + x: torch.Tensor, + t_pos: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + 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), + ) + positions = None if t_pos is None else t_pos.to(device=x.device, dtype=torch.float32) + 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, + t_pos=positions, + ) + 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, + t_pos=positions, + ) + return q, k, v + + +def full_with_keyframes( + x: torch.Tensor, + keyframe_x: torch.Tensor, + attn: NeighborhoodAttention3D, + norm: nn.RMSNorm, + scale: torch.Tensor, + shift: torch.Tensor, + keyframe_times: torch.Tensor, + keyframe_valid: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + if attn.joint_attention_function is None: + raise RuntimeError( + "keyframe decode needs joint_attention_function installed; build the decoder " + "through apply_diffvae_config / apply_diffvae_mode" + ) + batch, t, h, w, _ = x.shape + planes = keyframe_x.shape[1] + + y = norm(x) * (1.0 + scale) + shift + keyframe_y = norm(keyframe_x) * (1.0 + scale) + shift + + q, k, v = _qkv_nested_rope(attn, y) + keyframe_q, keyframe_k, keyframe_v = _qkv_nested_rope(attn, keyframe_y, keyframe_times) + q, k, v = q.contiguous(), k.contiguous(), v.contiguous() + keyframe_q = keyframe_q.contiguous() + keyframe_k = keyframe_k.contiguous() + keyframe_v = keyframe_v.contiguous() + + out, keyframe_out = attn.joint_attention_function( + attn, + q, + k, + v, + keyframe_q, + keyframe_k, + keyframe_v, + keyframe_times, + keyframe_valid, + ) + x = x + attn.proj(out.reshape(batch, t, h, w, attn.dim)) + keyframe_x = keyframe_x + attn.proj(keyframe_out.reshape(batch, planes, h, w, attn.dim)) + return x, keyframe_x + + +def full( + x: torch.Tensor, + attn: NeighborhoodAttention3D, + norm: nn.RMSNorm, + scale: torch.Tensor, + shift: torch.Tensor, +) -> torch.Tensor: + 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) + + +residual_attn = full +residual_attn_with_keyframes = full_with_keyframes + + +"""Combined pathway MLP: out-of-place AdaLN SwiGLU residual.""" + + +def residual_mlp( + x: torch.Tensor, + mlp: nn.Module, + norm: nn.RMSNorm, + scale: torch.Tensor, + shift: torch.Tensor, + tile: SwiGLUTileSpec, +) -> torch.Tensor: + y = modulate(norm(x), scale, shift) + if y.numel() == 0: + return x + return x + swiglu_tiled(y, *swiglu_weights(mlp), tile) + + +"""CombinedDiffusionNABlock: context_and_x inject + full-volume attn + residual MLP.""" + + +class CombinedDiffusionNABlock(DiffusionNABlock): + def forward_combined_with_keyframes( + self, + context_and_x: torch.Tensor, + keyframe_context_and_x: torch.Tensor, + modulation: tuple[torch.Tensor, ...], + keyframe_times: torch.Tensor, + keyframe_valid: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + scale_msa, shift_msa, scale_mlp, shift_mlp = self._modulation(modulation) + w_proj, b_proj = vram_ready_linear(self.context_proj) + x = inject_context(context_and_x, w_proj, b_proj) + keyframe_x = inject_context(keyframe_context_and_x, w_proj, b_proj) + x, keyframe_x = residual_attn_with_keyframes( + x, + keyframe_x, + self.attn, + self.norm1, + scale_msa, + shift_msa, + keyframe_times, + keyframe_valid, + ) + x = residual_mlp(x, self.mlp, self.norm2, scale_mlp, shift_mlp, self.mlp.tile) + keyframe_x = residual_mlp(keyframe_x, self.mlp, self.norm2, scale_mlp, shift_mlp, self.mlp.tile) + return x, keyframe_x * keyframe_valid[None, :, None, None, None] + + 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) + w_proj, b_proj = vram_ready_linear(self.context_proj) + x = inject_context(context_and_x, w_proj, b_proj) + 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) + + +"""Diffusion (NATTEN) video VAE decoder.""" + + +logger: logging.Logger = logging.getLogger(__name__) + + +_L_STAGE_CHANNELS: Tuple[int, ...] = (1024, 512, 256, 256, 128) + + +_L_STAGE_DEPTHS: Tuple[int, ...] = (4, 6, 4, 2, 2) + + +_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 +) + + +_L_STAGE_KERNELS: Tuple[Tuple[int, int, int], ...] = ( + (3, 7, 7), + (3, 7, 7), + (3, 5, 5), + (3, 5, 5), + (3, 3, 3), +) + + +_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, Disposable, VideoDecoder): + 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) + # Keyframe-stream tag, added to un-normalized keyframe latents before the shared + # ``conv_in`` and nowhere else. It is the only keyframe-specific weight in the + # whole feature. Checkpoints predating the keyframe training have no such key, so + # ``video_decoder_sd_ops_for_checkpoint`` synthesizes zeros -- a missing key would + # otherwise leave the parameter on the meta device under ``strict=False`` load. + self.type_emb = nn.Parameter(torch.zeros(in_channels)) + + 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 + 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] = 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] = compute_tile_min_size( + stage_kernels[3], stage5_kernel, up3_stride + ) + self.tile_halos: Tuple[Tuple[int, int, int], Tuple[int, int, int]] = 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 + # Remaining temporal upsampling per stage input, plus 1 for stage 5: the divisor in + # ``keyframe_stage_times``. (8, 8, 4, 2, 1) for the production ladder. + self._keyframe_time_strides: Tuple[int, ...] = remaining_time_strides(self.upsamples) + + def _run_det_stage(self, x: torch.Tensor, stage_i: int, drop_leading_frame: bool) -> torch.Tensor: + 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: + 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 _keyframe_stream_from_latents( + self, + keyframes: DecodeKeyframes, + *, + valid: torch.Tensor | None = None, + ) -> KeyframeStream: + latents = self.per_channel_statistics.un_normalize(keyframes.latents) + x = latents.permute(0, 2, 3, 4, 1) + x = x + self.type_emb.to(dtype=x.dtype, device=x.device).view(1, 1, 1, 1, -1) + x = self.conv_in(x) + planes = x.shape[1] + if valid is None: + valid = torch.ones(planes, dtype=torch.bool, device=x.device) + times = keyframe_clip_times( + keyframes.pixel_frame_indices, + self._keyframe_time_strides[0], + keyframes.clip_start_frame, + ) + return KeyframeStream(x=x, times=times.to(device=x.device), valid=valid.to(device=x.device)).masked() + + def _run_det_stage_with_keyframes( + self, + x: torch.Tensor, + keyframes: KeyframeStream, + stage_i: int, + drop_leading_frame: bool, + pixel_frame_indices: torch.Tensor, + next_time_origin: float, + clip_start_frame: int = 0, + ) -> tuple[torch.Tensor, KeyframeStream]: + 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, keyframes = block.forward_with_keyframes(x, keyframes) + x = self.upsamples[stage_i](x, drop_leading_frame=drop_leading_frame) + keyframe_x = upsample_keyframe_planes(self.upsamples[stage_i], keyframes.x) + next_times = keyframe_clip_times( + pixel_frame_indices, + self._keyframe_time_strides[stage_i + 1], + clip_start_frame, + extra_origin=next_time_origin, + ) + return x, KeyframeStream( + x=keyframe_x, + times=next_times.to(device=keyframe_x.device), + valid=keyframes.valid, + ).masked() + + def forward_stages_1_to_3_with_keyframes( + self, + z_noisy: torch.Tensor, + keyframes: DecodeKeyframes, + drop_leading_frame: bool = True, + *, + keyframe_valid: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, KeyframeStream]: + keyframes.validate() + if z_noisy.shape[-2:] != keyframes.latents.shape[-2:]: + raise ValueError( + f"keyframe latents must share the video latent's H/W (identical padding), got " + f"{tuple(keyframes.latents.shape[-2:])} vs {tuple(z_noisy.shape[-2:])}" + ) + stream = self._keyframe_stream_from_latents(keyframes, valid=keyframe_valid) + x = self.per_channel_statistics.un_normalize(z_noisy).permute(0, 2, 3, 4, 1) + x = self.conv_in(x) + for stage_i in range(3): + x, stream = self._run_det_stage_with_keyframes( + x, + stream, + stage_i, + drop_leading_frame, + keyframes.pixel_frame_indices, + 0.0, + clip_start_frame=keyframes.clip_start_frame, + ) + return x, stream + + def forward_stage_4_with_keyframes( + self, + x: torch.Tensor, + keyframes: KeyframeStream, + pixel_frame_indices: torch.Tensor, + drop_leading_frame: bool = True, + pad_trailing: bool = True, + *, + stage4_time_origin: float = 0.0, + pixel_time_origin: float = 0.0, + clip_start_frame: int = 0, + ) -> tuple[torch.Tensor, KeyframeStream]: + # Rebuild from global indices rather than trusting the caller's stream: stages 1-3 of a + # full-clip decode are global, and Dist has already folded clip_start into clip times. + keyframes = dataclasses.replace( + keyframes, + times=keyframe_clip_times( + pixel_frame_indices, + self._keyframe_time_strides[3], + clip_start_frame, + extra_origin=stage4_time_origin, + ).to(device=keyframes.x.device), + ) + if self.deferred_stage4_upsample: + return self._forward_stage_4_deferred_with_keyframes( + x, keyframes, pixel_frame_indices, pad_trailing, pixel_time_origin, clip_start_frame + ) + x, keyframes = self._run_det_stage_with_keyframes( + x, + keyframes, + 3, + drop_leading_frame, + pixel_frame_indices, + pixel_time_origin, + clip_start_frame=clip_start_frame, + ) + if pad_trailing: + x = 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, keyframes + + def _forward_stage_4_deferred_with_keyframes( + self, + x: torch.Tensor, + keyframes: KeyframeStream, + pixel_frame_indices: torch.Tensor, + pad_trailing: bool, + pixel_time_origin: float, + clip_start_frame: int = 0, + ) -> tuple[torch.Tensor, KeyframeStream]: + 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, keyframes = block.forward_with_keyframes(x, keyframes) + if pad_trailing: + up_t = int(self.upsamples[3].stride[0]) + x = 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)), + ) + stage5_times = keyframe_clip_times( + pixel_frame_indices, + self._keyframe_time_strides[4], + clip_start_frame, + extra_origin=pixel_time_origin, + ) + return x, KeyframeStream( + x=keyframes.x, + times=stage5_times.to(device=keyframes.x.device), + valid=keyframes.valid, + ).masked() + + def forward_stage_4( + self, + x: torch.Tensor, + drop_leading_frame: bool = True, + pad_trailing: bool = True, + ) -> torch.Tensor: + 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 = 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 = 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: + 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 _keyframe_context_and_x_for_diff_step( + self, + keyframe_context: torch.Tensor, + keyframe_x_t: torch.Tensor, + keyframe_valid: torch.Tensor, + ) -> torch.Tensor: + patched = patchify(keyframe_x_t, patch_size_hw=self.patch_size, patch_size_t=1) + x = self.conv_in_x_t(patched.permute(0, 2, 3, 4, 1)) + x = x * keyframe_valid[None, :, None, None, None] + return torch.cat([keyframe_context, x], dim=-1) + + def _x_for_diff_step(self, x_t: torch.Tensor) -> torch.Tensor: + 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 _keyframe_x_for_diff_step(self, keyframe_x_t: torch.Tensor, keyframe_valid: torch.Tensor) -> torch.Tensor: + patched = patchify(keyframe_x_t, patch_size_hw=self.patch_size, patch_size_t=1) + x = self.conv_in_x_t(patched.permute(0, 2, 3, 4, 1)) + return (x * keyframe_valid[None, :, None, None, None]).contiguous() + + def forward_diff_step( + self, + context_and_x: torch.Tensor, + t: torch.Tensor, + ) -> torch.Tensor: + 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)) + return self._pixels_from_stage5(x_half) + + def forward_diff_step_with_keyframes( + self, + context_and_x: torch.Tensor, + keyframe_context_and_x: torch.Tensor, + t: torch.Tensor, + keyframe_times: torch.Tensor, + keyframe_valid: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + x_half = context_and_x[..., self.context_channels :] + keyframe_half = keyframe_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) + # Keyframe dim 1 is the plane count, not T, and keyframe_times / keyframe_valid pin it. + # Marking it dynamic and then specializing to P raises ConstraintViolationError. + for dim in (2, 3): + torch._dynamo.mark_dynamic(keyframe_context_and_x, dim) + + for block in self.diff_blocks: + x_out, keyframe_out = block.forward_combined_with_keyframes( + context_and_x, + keyframe_context_and_x, + modulation, + keyframe_times, + keyframe_valid, + ) + x_half.copy_(x_out) + keyframe_half.copy_(keyframe_out) + + return self._pixels_from_stage5(x_half), self._pixels_from_stage5(keyframe_half) + + def _pixels_from_stage5(self, x: torch.Tensor) -> torch.Tensor: + 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: + 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) + + from ltx_core.model.video_vae.transformer.dsl_kernels import DSLDiffusionBlockChain # noqa: PLC0415 + + if isinstance(self.diff_blocks, DSLDiffusionBlockChain): + # Ping-pong fused launches; same deferred (x, stage4_feat) contract. + x = self.diff_blocks(x, stage4_feat, modulation, drop_leading_frame=drop_leading_frame) + else: + 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 forward_diff_step_deferred_with_keyframes( + self, + x: torch.Tensor, + stage4_feat: torch.Tensor, + keyframe_x: torch.Tensor, + keyframe_stage4_feat: torch.Tensor, + t: torch.Tensor, + keyframe_times: torch.Tensor, + keyframe_valid: torch.Tensor, + *, + drop_leading_frame: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + 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) + # Plane count, not T -- see the note in the combined path above. + for dim in (2, 3): + torch._dynamo.mark_dynamic(keyframe_x, dim) + torch._dynamo.mark_dynamic(keyframe_stage4_feat, dim) + + # The DSL chain drives itself, so both streams recycle output buffers instead of + # allocating a volume per block per stream; the chunked blocks expose the same + # ``forward_x_ctx_with_keyframes``, so the fallback loop drives those. + from ltx_core.model.video_vae.transformer.dsl_kernels import DSLDiffusionBlockChain # noqa: PLC0415 + + if isinstance(self.diff_blocks, DSLDiffusionBlockChain): + x, keyframe_x = self.diff_blocks.forward_x_ctx_with_keyframes( + x, + stage4_feat, + keyframe_x, + keyframe_stage4_feat, + modulation, + keyframe_times, + keyframe_valid, + drop_leading_frame=drop_leading_frame, + ) + else: + for block in self.diff_blocks: + x, keyframe_x = block.forward_x_ctx_with_keyframes( + x, + stage4_feat, + keyframe_x, + keyframe_stage4_feat, + modulation, + keyframe_times, + keyframe_valid, + drop_leading_frame=drop_leading_frame, + ) + + return self._pixels_from_stage5(x), self._pixels_from_stage5(keyframe_x) + + def _euler_step( + self, x_t: torch.Tensor, model_out: torch.Tensor, t_now: torch.Tensor, t_next: torch.Tensor + ) -> torch.Tensor: + 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: + 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 _stage5_canvas_from_context( + self, + context_tile: torch.Tensor, + *, + drop_leading_frame: bool, + ) -> tuple[int, int, int]: + t, h, w = context_tile.shape[1], context_tile.shape[2], context_tile.shape[3] + if self.deferred_stage4_upsample: + # Context is still pre-upsample, so this is the same geometry as a stage-4 + # input. Ghost crop already ran, so do not re-apply the kernel-T floor. + return stage5_pixel_shape_from_stage4( + t, + h, + w, + upsample_stride=tuple(self.upsamples[3].stride), # type: ignore[arg-type] + patch_size=self.patch_size, + stage5_kernel_t=self.stage5_kernel[0], + drop_leading_frame=drop_leading_frame, + pad_trailing=False, + ) + return t, h * self.patch_size, w * self.patch_size + + def _decode_one_tile_with_keyframes( # noqa: PLR0913 + self, + feat_tile: torch.Tensor, + keyframes: KeyframeStream, + pixel_frame_indices: torch.Tensor, + *, + is_origin: bool, + timestep: torch.Tensor, + pad_trailing: bool, + generator: torch.Generator | None, + compute_dtype: torch.dtype, + x_t_tile_init: torch.Tensor | None = None, + stage4_time_origin: float = 0.0, + pixel_time_origin: float = 0.0, + clip_start_frame: int = 0, + ) -> torch.Tensor: + context_tile, keyframes = self.forward_stage_4_with_keyframes( + feat_tile, + keyframes, + pixel_frame_indices, + drop_leading_frame=is_origin, + pad_trailing=pad_trailing, + stage4_time_origin=stage4_time_origin, + pixel_time_origin=pixel_time_origin, + clip_start_frame=clip_start_frame, + ) + + batch = context_tile.shape[0] + canvas_t, canvas_h, canvas_w = self._stage5_canvas_from_context(context_tile, drop_leading_frame=is_origin) + randn_device = generator.device if generator is not None else feat_tile.device + + def _noise(frames: int) -> torch.Tensor: + return torch.randn( + (batch, self.out_channels, frames, canvas_h, canvas_w), + dtype=compute_dtype, + generator=generator, + device=randn_device, + ).to(feat_tile.device) + + x_t = _noise(canvas_t) if x_t_tile_init is None else x_t_tile_init + keyframe_x_t = _noise(keyframes.num_planes) + + def _step(t_now: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + if self.deferred_stage4_upsample: + return self.forward_diff_step_deferred_with_keyframes( + self._x_for_diff_step(x_t), + context_tile, + self._keyframe_x_for_diff_step(keyframe_x_t, keyframes.valid), + keyframes.x, + t_now, + keyframes.times, + keyframes.valid, + drop_leading_frame=is_origin, + ) + context_and_x = self._context_and_x_for_diff_step(context_tile, x_t) + keyframe_context_and_x = self._keyframe_context_and_x_for_diff_step( + keyframes.x, keyframe_x_t, keyframes.valid + ) + return self.forward_diff_step_with_keyframes( + context_and_x, keyframe_context_and_x, t_now, keyframes.times, keyframes.valid + ) + + _, num_steps = timestep.shape + for i in range(num_steps - 1): + t_now = timestep[:, i] + t_next = timestep[:, i + 1] + video_out, keyframe_out = _step(t_now) + x_t = self._euler_step(x_t, video_out.to(torch.float32), t_now, t_next) + keyframe_x_t = self._euler_step(keyframe_x_t, keyframe_out.to(torch.float32), t_now, t_next) + + t_now = timestep[:, -1] + video_out, _ = _step(t_now) + if self.model_output_type == "x0": + return video_out + return self._euler_step(x_t, video_out.to(torch.float32), t_now, torch.zeros_like(t_now)) + + def _decode_temporal_group_isolated_with_keyframes( # noqa: PLR0913 + self, + tiles: List[Tile], + feat_s4: torch.Tensor, + stream: KeyframeStream, + pixel_frame_indices: 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, + clip_start_frame: int = 0, + ) -> Tuple[torch.Tensor, torch.Tensor | None]: + 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 + up3_stride = tuple(self.upsamples[3].stride) + + for tile_index, tile in tqdm(enumerate(tiles), total=len(tiles), desc="DiffVAE decode"): + feat_tile, is_origin, pad_trailing, content_thw = slice_stage4_tile( + feat_s4, tile, content_frames=content_s4_frames + ) + # Two origins at two scales -- see forward_stage_4_with_keyframes. + stage4_origin = tile.in_coords[1].indices(content_s4_frames)[0] + pixel_lo, pixel_hi, _ = tile.out_coords[2].indices(full_torch_shape[2]) + + # ``out_coords`` are local to this latent. Dist slices keep global indices and + # ``clip_start_frame`` as the origin, so a local ``[0, 72)`` still has to select + # global ``[56, 127]``. + keep = planes_for_tile(pixel_frame_indices, pixel_lo, pixel_hi - 1, clip_start_frame=clip_start_frame) + if not bool(keep.any()): + raise RuntimeError( + f"tile covering pixel frames [{pixel_lo + clip_start_frame}, " + f"{pixel_hi - 1 + clip_start_frame}] selected no keyframe planes " + f"out of {int(pixel_frame_indices.shape[0])}; planes_for_tile always keeps at least one" + ) + tile_stream = stream.select_planes(keep.to(stream.valid.device)).crop_spatial( + tile.in_coords[2], tile.in_coords[3] + ) + # Not debug-only: the decode below needs this tile's plane positions. + tile_indices = pixel_frame_indices[keep.to(pixel_frame_indices.device)] + if logger.isEnabledFor(logging.INFO): + logger.info( + "keyframe decode: tile %d/%d frames [%d, %d), stage-4 extent %dx%dx%d, %d of %d planes at %s", + tile_index + 1, + len(tiles), + pixel_lo, + pixel_hi, + feat_tile.shape[1], + feat_tile.shape[2], + feat_tile.shape[3], + tile_stream.num_planes, + int(pixel_frame_indices.shape[0]), + tile_indices.tolist(), + ) + + x_t_tile_init: torch.Tensor | None = None + if x_t_init is not None: + stage5_f, stage5_h, stage5_w = 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, + ) + # Same edge policy as the plain path: expand/crop the shared noise field + # rather than drawing fresh noise, since NA mixes padded values inward. + x_t_tile_init = x_t_init[tile.out_coords] + x_t_tile_init, _ = resize_axis(x_t_tile_init, 2, stage5_f, mode="repeat_last") + x_t_tile_init, _ = resize_axis(x_t_tile_init, 3, stage5_h, mode="symmetric") + x_t_tile_init, _ = resize_axis(x_t_tile_init, 4, stage5_w, mode="symmetric") + + pixel_tile = self._decode_one_tile_with_keyframes( + feat_tile, + tile_stream, + tile_indices, + is_origin=is_origin, + timestep=timestep, + pad_trailing=pad_trailing, + generator=generator, + compute_dtype=compute_dtype, + x_t_tile_init=x_t_tile_init, + stage4_time_origin=float(stage4_origin), + pixel_time_origin=float(pixel_lo), + clip_start_frame=clip_start_frame, + ) + content_pixel_shape = pixel_tile_shape(full_torch_shape, tile.out_coords) + pixel_tile = 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_groups_with_keyframes( # noqa: PLR0913, PLR0915 + self, + feat_s4: torch.Tensor, + stream: KeyframeStream, + pixel_frame_indices: torch.Tensor, + latent: torch.Tensor, + tiling_config: TilingConfig, + timestep: torch.Tensor, + generator: torch.Generator | None, + *, + content_pixel: VideoLatentShape, + h_pad: AxisPad | None, + w_pad: AxisPad | None, + as_fhwc: bool, + clip_start_frame: int = 0, + ) -> Iterator[torch.Tensor]: + full_video_shape = ( + VideoLatentShape.from_torch_shape(latent.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 = stage4_thw_from_latent( + strides, latent.shape[2], latent.shape[3], latent.shape[4], drop_leading_frame=True + ) + tiles = 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, + ) + 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] + + 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: + randn_device = generator.device if generator is not None else latent.device + x_t_init = torch.randn( + tuple(target_shape), dtype=latent.dtype, generator=generator, device=randn_device + ).to(latent.device) + + logger.info( + "keyframe decode: %d tile(s) in %d temporal group(s), %d frames at %dx%d, %d planes", + len(tiles), + len(groups), + content_pixel.frames, + content_pixel.height, + content_pixel.width, + int(pixel_frame_indices.shape[0]), + ) + + scaled_h_pad = scale_axis_pad(h_pad, self.video_downscale_factors.height) + scaled_w_pad = scale_axis_pad(w_pad, self.video_downscale_factors.width) + overlap_stub: torch.Tensor | None = None + overlap_stub_weights: torch.Tensor | None = None + + def _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 + chunk = buf[:, :, :frames_keep] + if wts is not None: + floor = _weight_floor(wts.dtype) + chunk = chunk / wts[:, :, :frames_keep].clamp(min=floor) + chunk = crop_pixels_to_content( + chunk.to(latent.dtype), + frames_keep, + content_pixel.height, + content_pixel.width, + h_pad=scaled_h_pad, + w_pad=scaled_w_pad, + ) + return chunk[0].permute(1, 2, 3, 0).contiguous() if as_fhwc else chunk + + for group_index, group in enumerate(groups): + curr_temporal_slice = group_slices[group_index] + logger.info( + "keyframe decode: group %d/%d, frames [%d, %d)", + group_index + 1, + len(groups), + curr_temporal_slice.start, + curr_temporal_slice.stop, + ) + buffer, weights = self._decode_temporal_group_isolated_with_keyframes( + group, + feat_s4, + stream, + pixel_frame_indices, + s4_t, + x_t_init, + timestep, + full_video_shape, + curr_temporal_slice, + generator, + complementary=complementary, + clip_start_frame=clip_start_frame, + ) + + if overlap_stub is not None: + overlap_len = int(overlap_stub.shape[2]) + if overlap_len > 0: + overlap_stub += buffer[:, :, :overlap_len] + buffer[:, :, :overlap_len] = overlap_stub + if not complementary: + assert overlap_stub_weights is not None + assert weights is not None + overlap_stub_weights += weights[:, :, :overlap_len] + weights[:, :, :overlap_len] = overlap_stub_weights + overlap_stub = None + overlap_stub_weights = None + + if group_index + 1 < len(groups): + next_start = group_slices[group_index + 1].start + exclusive_len = min(max(0, next_start - curr_temporal_slice.start), buffer.shape[2]) + emitted = _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 group's 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 = _emit(buffer, weights, curr_temporal_slice.start) + if emitted is not None: + yield emitted + + def _decode_pixels_with_keyframes( + self, + latent: torch.Tensor, + keyframes: DecodeKeyframes, + tiling_config: TilingConfig | None = None, + generator: torch.Generator | None = None, + *, + as_fhwc: bool = False, + ) -> Iterator[torch.Tensor]: + content_shape = VideoLatentShape.from_torch_shape(latent.shape) + content_pixel = content_shape.upscale(self.video_downscale_factors)._replace(channels=self.out_channels) + keyframes.validate(num_frames=content_pixel.frames) + + latent, (_t_pad, h_pad, w_pad) = ensure_min_latent_shape(latent, self.stage_min_tile_sizes) + # Same spatial floor for the keyframe planes, with the plane axis pinned by a + # temporal minimum of 1. The pad is symmetric, so padding only one stream would + # offset every plane from the video by half of it. + _min_t, min_h, min_w = self.stage_min_tile_sizes + keyframe_latents, (_, keyframe_h_pad, keyframe_w_pad) = ensure_min_latent_shape( + keyframes.latents, (1, min_h, min_w) + ) + if (keyframe_h_pad, keyframe_w_pad) != (h_pad, w_pad): + raise RuntimeError( + f"keyframe spatial pad {(keyframe_h_pad, keyframe_w_pad)} != video pad {(h_pad, w_pad)}; " + "the two streams must share one spatial origin" + ) + padded_keyframes = dataclasses.replace(keyframes, latents=keyframe_latents) + + # Ghost pad is a temporal-border workaround for the video stream; keyframe planes have + # no temporal extent to pad. The appendix is cropped off context before stage 5. + latent_padded = pad_trailing_latent_for_natten_border( + latent, self._natten_trailing_pad_latent_frames + ) + feat_s4, stream = self.forward_stages_1_to_3_with_keyframes( + latent_padded, padded_keyframes, drop_leading_frame=True + ) + + batch = latent.shape[0] + timestep = self.default_inference_timesteps.to(latent.device).unsqueeze(0).expand(batch, -1) + if tiling_config is not None: + yield from self._decode_groups_with_keyframes( + feat_s4, + stream, + keyframes.pixel_frame_indices, + latent, + tiling_config, + timestep, + generator, + content_pixel=content_pixel, + h_pad=h_pad, + w_pad=w_pad, + as_fhwc=as_fhwc, + clip_start_frame=keyframes.clip_start_frame, + ) + return + + logger.info("keyframe decode: untiled, %d frames, %d planes", content_pixel.frames, stream.num_planes) + pixels = self._decode_one_tile_with_keyframes( + feat_s4, + stream, + keyframes.pixel_frame_indices, + is_origin=True, + timestep=timestep, + pad_trailing=True, + generator=generator, + compute_dtype=latent.dtype, + clip_start_frame=keyframes.clip_start_frame, + ) + pixels = crop_pixels_to_content( + pixels, + content_pixel.frames, + content_pixel.height, + content_pixel.width, + h_pad=scale_axis_pad(h_pad, self.video_downscale_factors.height), + w_pad=scale_axis_pad(w_pad, self.video_downscale_factors.width), + ).to(latent.dtype) + if as_fhwc: + yield pixels[0].permute(1, 2, 3, 0).contiguous() + else: + yield pixels + + def _decode_video_with_keyframes( + self, + latent: torch.Tensor, + keyframes: DecodeKeyframes, + tiling_config: TilingConfig | None = None, + generator: torch.Generator | None = None, + ) -> Iterator[torch.Tensor]: + def to_rgb(frames: torch.Tensor) -> torch.Tensor: + return frames.add_(1).mul_(0.5).clamp_(0, 1) + + for chunk in self._decode_pixels_with_keyframes( + latent, keyframes, tiling_config, generator=generator, as_fhwc=True + ): + yield to_rgb(chunk) + + 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]: + 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 tqdm(tiles, total=len(tiles), desc="DiffVAE decode"): + feat_tile, is_origin, pad_trailing, content_thw = slice_stage4_tile( + feat_s4, tile, content_frames=content_s4_frames + ) + content_pixel_shape = pixel_tile_shape(full_torch_shape, tile.out_coords) + stage5_f, stage5_h, stage5_w = 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, _ = resize_axis(x_t_tile_init, 2, stage5_f, mode="repeat_last") + x_t_tile_init, _ = resize_axis(x_t_tile_init, 3, stage5_h, mode="symmetric") + x_t_tile_init, _ = 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 = 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]: + 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) = 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 = stage4_thw_from_latent( + strides, latent.shape[2], latent.shape[3], latent.shape[4], drop_leading_frame=True + ) + tiles = 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 = 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=_weight_floor(wts.dtype)) + return (buf / wts).to(latent.dtype) + + def _narrow_content_cfhw(t: torch.Tensor, frames_keep: int) -> torch.Tensor: + 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 = 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, _ = resize_axis(x, 3, th, mode="symmetric") + if w_pad is not None: + before = 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, _ = 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 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_(_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: + return next(self._decode_pixels(sample, tiling_config=None, generator=generator)) + + def decode_video( + self, + latent: torch.Tensor, + tiling_config: TilingConfig | None = None, + generator: torch.Generator | None = None, + *, + keyframes: DecodeKeyframes | None = None, + ) -> Iterator[torch.Tensor]: + if keyframes is not None: + yield from self._decode_video_with_keyframes(latent, keyframes, tiling_config, generator) + return + + 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) + +class LTX25DiffusionVideoDecoder(DiffusionVideoDecoder): + def forward(self, sample, generator=None, keyframes=None): + return self.decode(sample, generator=generator, keyframes=keyframes) + + def auto_tiling_config(self, latent, keyframes=None): + pixel_shape = ( + VideoLatentShape.from_torch_shape(latent.shape) + .upscale(self.video_downscale_factors) + ._replace(channels=self.out_channels) + ) + upsample_strides = [tuple(upsample.stride) for upsample in self.upsamples] + pixel_scale = stage4_to_pixel_scale_factors(upsample_strides[3], self.patch_size) + overlap_t, overlap_hw = recommended_pixel_overlaps(self.tile_halos, pixel_scale) + ft, fh, fw = pixel_scale.time, pixel_scale.height, pixel_scale.width + 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 = _round_up(max(2 * ft, 2 * overlap_t, _round_up(self.tile_min_sizes[0] * ft, ft), 16), step_t) + min_h = _round_up(max(2 * fh, 2 * overlap_hw, _round_up(self.tile_min_sizes[1] * fh, fh), 512), step_h) + min_w = _round_up(max(2 * fw, 2 * overlap_hw, _round_up(self.tile_min_sizes[2] * fw, fw), 512), step_w) + return TileSizeConfig( + frames=DimensionSizeConfig(min_t, overlap_t), + height=DimensionSizeConfig(min_h, overlap_hw), + width=DimensionSizeConfig(min_w, overlap_hw), + ) + + def _resolve_tiling_config(self, latent, tiled, keyframes=None, tile_size_in_pixels=None, tile_size_in_frames=None): + if not tiled: + return None + if tile_size_in_pixels is not None and tile_size_in_frames is not None: + try: + pixel_scale = stage4_to_pixel_scale_factors([tuple(upsample.stride) for upsample in self.upsamples][3], self.patch_size) + overlap_t, overlap_hw = recommended_pixel_overlaps(self.tile_halos, pixel_scale) + tiling_config = TileSizeConfig( + frames=DimensionSizeConfig(tile_size_in_frames, overlap_t), + height=DimensionSizeConfig(tile_size_in_pixels, overlap_hw), + width=DimensionSizeConfig(tile_size_in_pixels, overlap_hw), + ) + pixel_shape = ( + VideoLatentShape.from_torch_shape(latent.shape) + .upscale(self.video_downscale_factors) + ._replace(channels=self.out_channels) + ) + tiling_config.validate(pixel_scale, pixel_shape) + return tiling_config + except ValueError: + pass + return self.auto_tiling_config(latent, keyframes=keyframes) + + def decode( + self, + latent, + tiled=False, + seed=None, + rand_device="cpu", + keyframes=None, + tile_size_in_pixels=None, + tile_size_in_frames=None, + **kwargs, + ): + generator = torch.Generator(device=rand_device).manual_seed(42 if seed is None else seed + 42) + tiling_config = self._resolve_tiling_config(latent, tiled, keyframes, tile_size_in_pixels, tile_size_in_frames) + iterator = ( + self._decode_pixels_with_keyframes(latent, keyframes, tiling_config, generator=generator) + if keyframes is not None + else self._decode_pixels(latent, tiling_config, generator=generator) + ) + chunks = list(iterator) + if not chunks: + raise RuntimeError("Diffusion decoder produced no output chunks") + return torch.cat(chunks, dim=2) diff --git a/diffsynth/models/ltx25_duration_head.py b/diffsynth/models/ltx25_duration_head.py new file mode 100644 index 000000000..9fdf01ccc --- /dev/null +++ b/diffsynth/models/ltx25_duration_head.py @@ -0,0 +1,57 @@ +import torch +from torch import nn + + +class LTX25AttentionPooler(nn.Module): + def __init__(self, hidden_dim: int = 256, num_queries: int = 1, num_heads: int = 4): + super().__init__() + self.hidden_dim = hidden_dim + self.num_queries = num_queries + self.query_tokens = nn.Parameter(torch.randn(num_queries, hidden_dim) * 0.02) + self.cross_attn = nn.MultiheadAttention( + embed_dim=hidden_dim, + num_heads=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): + def __init__( + self, + video_cross_attention_dim: int = 4096, + audio_cross_attention_dim: int = 2048, + pooler_hidden_dim: int = 256, + num_queries: int = 1, + num_pooler_heads: int = 4, + mlp_hidden: int = 256, + ): + super().__init__() + self.pooler_hidden_dim = pooler_hidden_dim + 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: + if video_tokens is None and audio_tokens is None: + raise ValueError("LTX25DurationHead.forward requires video_tokens and/or audio_tokens.") + token_groups = [] + if video_tokens is not None: + token_groups.append(self.video_input_proj(video_tokens) + self.video_modality_emb) + if audio_tokens is not None: + token_groups.append(self.audio_input_proj(audio_tokens) + self.audio_modality_emb) + pooled = self.attention_pooler(torch.cat(token_groups, dim=1)) + hidden = torch.nn.functional.gelu(self.mlp_hidden(pooled.reshape(pooled.shape[0], -1)), approximate="tanh") + return self.mlp_out(hidden).squeeze(-1).exp() diff --git a/diffsynth/models/ltx25_text_encoder.py b/diffsynth/models/ltx25_text_encoder.py new file mode 100644 index 000000000..f0d05a0af --- /dev/null +++ b/diffsynth/models/ltx25_text_encoder.py @@ -0,0 +1,449 @@ +import math +from pathlib import Path +from typing import NamedTuple + +import torch +from transformers import Gemma4UnifiedConfig, Gemma4UnifiedForConditionalGeneration, PreTrainedTokenizerFast + +from .ltx2_common import rms_norm +from .ltx2_dit import ( + Attention, + FeedForward, + LTXRopeType, + generate_freq_grid_np, + generate_freq_grid_pytorch, + precompute_freqs_cis, +) + + +class LTX25TextEncoder(Gemma4UnifiedForConditionalGeneration): + def __init__(self): + config = { + "architectures": ["Gemma4UnifiedForConditionalGeneration"], + "audio_config": { + "_name_or_path": "", + "architectures": None, + "audio_embed_dim": 640, + "chunk_size_feed_forward": 0, + "dtype": "bfloat16", + "id2label": {"0": "LABEL_0", "1": "LABEL_1"}, + "initializer_range": 0.02, + "is_encoder_decoder": False, + "label2id": {"LABEL_0": 0, "LABEL_1": 1}, + "model_type": "gemma4_unified_audio", + "output_attentions": False, + "output_hidden_states": False, + "problem_type": None, + "return_dict": True, + "rms_norm_eps": 1e-06, + }, + "audio_token_id": 258881, + "boa_token_id": 256000, + "boi_token_id": 255999, + "dtype": "bfloat16", + "eoa_token_index": 258883, + "eoi_token_id": 258882, + "eos_token_id": [1, 106], + "gemma_version": "gemma4-12b-ltx-v1", + "image_token_id": 258880, + "initializer_range": 0.02, + "model_type": "gemma4_unified", + "text_config": { + "attention_bias": False, + "attention_dropout": 0.0, + "attention_k_eq_v": True, + "bos_token_id": 2, + "dtype": "bfloat16", + "enable_moe_block": False, + "eos_token_id": 1, + "final_logit_softcapping": 30.0, + "global_head_dim": 512, + "head_dim": 256, + "hidden_activation": "gelu_pytorch_tanh", + "hidden_size": 3840, + "hidden_size_per_layer_input": 0, + "initializer_range": 0.02, + "intermediate_size": 15360, + "layer_types": (["sliding_attention"] * 5 + ["full_attention"]) * 8, + "max_position_embeddings": 262144, + "model_type": "gemma4_unified_text", + "moe_intermediate_size": None, + "num_attention_heads": 16, + "num_experts": None, + "num_global_key_value_heads": 1, + "num_hidden_layers": 48, + "num_key_value_heads": 8, + "num_kv_shared_layers": 0, + "pad_token_id": 0, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "full_attention": {"partial_rotary_factor": 0.25, "rope_theta": 1000000.0, "rope_type": "proportional"}, + "sliding_attention": {"rope_theta": 10000.0, "rope_type": "default"}, + }, + "sliding_window": 1024, + "tie_word_embeddings": True, + "top_k_experts": None, + "use_bidirectional_attention": "vision", + "use_cache": True, + "use_double_wide_mlp": False, + "vocab_size": 262144, + "vocab_size_per_layer_input": 262144, + }, + "tie_word_embeddings": True, + "transformers_version": "5.10.1", + "video_token_id": 258884, + "vision_config": { + "_name_or_path": "", + "architectures": None, + "chunk_size_feed_forward": 0, + "dtype": "bfloat16", + "id2label": {"0": "LABEL_0", "1": "LABEL_1"}, + "initializer_range": 0.02, + "is_encoder_decoder": False, + "label2id": {"LABEL_0": 0, "LABEL_1": 1}, + "mm_embed_dim": 3840, + "mm_posemb_size": 1120, + "model_type": "gemma4_unified_vision", + "num_soft_tokens": 280, + "output_attentions": False, + "output_hidden_states": False, + "output_proj_dims": 3840, + "patch_size": 16, + "pooling_kernel_size": 3, + "problem_type": None, + "return_dict": True, + "rms_norm_eps": 1e-06, + }, + } + super().__init__(Gemma4UnifiedConfig(**config)) + # Gemma4 registers a constant `layer_scalar` as a persistent buffer on each decoder layer. + # Disk-offload reload restores only named_parameters(), so a persistent buffer is never fetched + # from disk and strict load_state_dict fails; register it as a frozen parameter instead. + for module in self.modules(): + if "layer_scalar" in module._buffers: + module.register_parameter("layer_scalar", torch.nn.Parameter(module._buffers.pop("layer_scalar"), requires_grad=False)) + + +class LTX25GemmaTokenizer: + def __init__(self, tokenizer_path: str | Path, max_length: int = 1024): + self.tokenizer = PreTrainedTokenizerFast.from_pretrained(tokenizer_path, local_files_only=True, model_max_length=max_length) + self.tokenizer.model_max_length = max_length + self.tokenizer.padding_side = "left" + if self.tokenizer.pad_token is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + self.max_length = max_length + + def tokenize_with_weights(self, text: str) -> dict[str, list[tuple[int, int]]]: + text = text.strip() + bos_id = self.tokenizer.bos_token_id + if bos_id is None: + raise ValueError("Packed Gemma tokenizer has no BOS token id.") + encoded = self.tokenizer(text, padding=False, truncation=True, max_length=self.max_length, return_tensors="pt") + input_ids = encoded.input_ids[0].tolist() + if not input_ids or input_ids[0] != bos_id: + input_ids = [bos_id, *input_ids][: self.max_length] + padded = self.tokenizer.pad( + {"input_ids": [input_ids]}, + padding="max_length", + max_length=self.max_length, + return_tensors="pt", + return_attention_mask=True, + ) + return {"gemma": list(zip(padded.input_ids[0].tolist(), padded.attention_mask[0].tolist(), strict=True))} + + +def norm_and_concat_per_token_rms( + encoded_text: torch.Tensor, + attention_mask: torch.Tensor, +) -> torch.Tensor: + batch_size, sequence_length, embedding_dim, num_layers = encoded_text.shape + variance = torch.mean(encoded_text**2, dim=2, keepdim=True) + normed = encoded_text * torch.rsqrt(variance + 1e-6) + normed = normed.reshape(batch_size, sequence_length, embedding_dim * num_layers) + return torch.where(attention_mask.bool().unsqueeze(-1), normed, torch.zeros_like(normed)) + + +def _rescale_norm(x: torch.Tensor, target_dim: int, source_dim: int) -> torch.Tensor: + return x * math.sqrt(target_dim / source_dim) + + +class LTX25FeatureExtractorV2(torch.nn.Module): + def __init__( + self, + embedding_dim: int = 3840, + num_layers: int = 49, + video_out_features: int = 4096, + audio_out_features: int = 2048, + ): + super().__init__() + self.video_aggregate_embed = torch.nn.Linear(embedding_dim * num_layers, video_out_features, bias=True) + self.audio_aggregate_embed = ( + torch.nn.Linear(embedding_dim * num_layers, audio_out_features, bias=True) + if audio_out_features is not None + else None + ) + self.embedding_dim = embedding_dim + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + padding_side: str = "left", + ) -> tuple[torch.Tensor, torch.Tensor | None]: + del padding_side + encoded = torch.stack(hidden_states, dim=-1) if isinstance(hidden_states, (list, tuple)) else hidden_states + normed = norm_and_concat_per_token_rms(encoded, attention_mask).to(encoded.dtype) + video = self.video_aggregate_embed( + _rescale_norm(normed, self.video_aggregate_embed.out_features, self.embedding_dim) + ) + audio = None + if self.audio_aggregate_embed is not None: + audio = self.audio_aggregate_embed( + _rescale_norm(normed, self.audio_aggregate_embed.out_features, self.embedding_dim) + ) + return video, audio + + +class LTX25Embeddings1DConnector(torch.nn.Module): + def __init__( + self, + attention_head_dim: int, + num_attention_heads: int, + num_layers: int, + positional_embedding_theta: float = 10000.0, + positional_embedding_max_pos: list[int] | None = None, + num_learnable_registers: int | None = 128, + rope_type: LTXRopeType = LTXRopeType.SPLIT, + double_precision_rope: bool = True, + apply_gated_attention: bool = True, + ff_bias: bool = False, + ): + super().__init__() + self.num_attention_heads = num_attention_heads + self.inner_dim = num_attention_heads * attention_head_dim + self.positional_embedding_theta = positional_embedding_theta + self.positional_embedding_max_pos = positional_embedding_max_pos if positional_embedding_max_pos is not None else [1] + self.rope_type = rope_type + self.double_precision_rope = double_precision_rope + self.transformer_1d_blocks = torch.nn.ModuleList( + [ + LTX25BasicTransformerBlock1D( + dim=self.inner_dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + rope_type=rope_type, + apply_gated_attention=apply_gated_attention, + ff_bias=ff_bias, + ) + for _ in range(num_layers) + ] + ) + self.num_learnable_registers = num_learnable_registers + if self.num_learnable_registers: + self.learnable_registers = torch.nn.Parameter( + torch.rand(self.num_learnable_registers, self.inner_dim, dtype=torch.bfloat16) * 2.0 - 1.0 + ) + + def _replace_padded_with_learnable_registers( + self, + hidden_states: torch.Tensor, + additive_attention_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, sequence_length, _ = hidden_states.shape + assert sequence_length % self.num_learnable_registers == 0 + registers = self.learnable_registers.to(hidden_states).repeat(sequence_length // self.num_learnable_registers, 1) + registers = registers.unsqueeze(0).expand(batch_size, -1, -1) + binary_mask = (additive_attention_mask[:, 0, 0, :].unsqueeze(-1) >= 0).to(hidden_states.dtype) + hidden_states = binary_mask * hidden_states + (1 - binary_mask) * registers + return hidden_states, torch.zeros_like(additive_attention_mask) + + def forward( + self, + hidden_states: torch.Tensor, + additive_attention_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.num_learnable_registers: + hidden_states, additive_attention_mask = self._replace_padded_with_learnable_registers( + hidden_states, + additive_attention_mask, + ) + indices_grid = torch.arange(hidden_states.shape[1], dtype=torch.float32, device=hidden_states.device) + indices_grid = indices_grid[None, None, :].expand(hidden_states.shape[0], -1, -1) + freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch + freqs_cis = precompute_freqs_cis( + indices_grid=indices_grid, + dim=self.inner_dim, + out_dtype=hidden_states.dtype, + theta=self.positional_embedding_theta, + max_pos=self.positional_embedding_max_pos, + num_attention_heads=self.num_attention_heads, + rope_type=self.rope_type, + freq_grid_generator=freq_grid_generator, + ) + for block in self.transformer_1d_blocks: + hidden_states = block(hidden_states, additive_attention_mask=additive_attention_mask, pe=freqs_cis) + return rms_norm(hidden_states), additive_attention_mask + + +class LTX25BasicTransformerBlock1D(torch.nn.Module): + def __init__( + self, + dim: int, + heads: int, + dim_head: int, + rope_type: LTXRopeType, + apply_gated_attention: bool, + ff_bias: bool, + ): + super().__init__() + self.attn1 = Attention( + query_dim=dim, + heads=heads, + dim_head=dim_head, + rope_type=rope_type, + apply_gated_attention=apply_gated_attention, + ) + self.ff = FeedForward(dim, dim_out=dim, bias=ff_bias) + + def forward( + self, + hidden_states: torch.Tensor, + additive_attention_mask: torch.Tensor | None = None, + pe: torch.Tensor | None = None, + ) -> torch.Tensor: + norm_hidden_states = rms_norm(hidden_states).squeeze(1) + hidden_states = self.attn1(norm_hidden_states, mask=additive_attention_mask, pe=pe) + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + hidden_states = self.ff(rms_norm(hidden_states)) + hidden_states + return hidden_states.squeeze(1) if hidden_states.ndim == 4 else hidden_states + + +class LTX25EmbeddingsProcessorOutput(NamedTuple): + video_encoding: torch.Tensor + audio_encoding: torch.Tensor | None + attention_mask: torch.Tensor + + +def _convert_to_additive_mask(attention_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + return (attention_mask.to(torch.int64) - 1).to(dtype).reshape( + attention_mask.shape[0], 1, 1, attention_mask.shape[-1] + ) * torch.finfo(dtype).max + + +def _right_pad_order(additive_attention_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + binary = (additive_attention_mask[:, 0, 0, :] >= 0).to(torch.int32) + sort_indices = torch.argsort(binary, dim=-1, descending=True, stable=True) + reordered = torch.gather(binary, 1, sort_indices) + additive = (reordered.to(additive_attention_mask.dtype) - 1) * torch.finfo(additive_attention_mask.dtype).max + return sort_indices, additive[:, None, None, :] + + +class LTX25EmbeddingsConnectors(torch.nn.Module): + def __init__( + self, + video_attention_heads: int = 32, + video_attention_head_dim: int = 128, + audio_attention_heads: int = 32, + audio_attention_head_dim: int = 64, + num_connector_layers: int = 8, + connector_max_positions: list[int] | None = None, + connector_ff_bias: bool = True, + ): + super().__init__() + connector_max_positions = [4096] if connector_max_positions is None else connector_max_positions + self.video_connector = LTX25Embeddings1DConnector( + attention_head_dim=video_attention_head_dim, + num_attention_heads=video_attention_heads, + num_layers=num_connector_layers, + positional_embedding_max_pos=connector_max_positions, + ff_bias=connector_ff_bias, + ) + self.audio_connector = LTX25Embeddings1DConnector( + attention_head_dim=audio_attention_head_dim, + num_attention_heads=audio_attention_heads, + num_layers=num_connector_layers, + positional_embedding_max_pos=connector_max_positions, + ff_bias=connector_ff_bias, + ) + + +class LTX25TextEncoderPostModules(torch.nn.Module): + def __init__( + self, + embedding_dim: int = 3840, + num_layers: int = 49, + video_attention_heads: int = 32, + video_attention_head_dim: int = 128, + audio_attention_heads: int = 32, + audio_attention_head_dim: int = 64, + num_connector_layers: int = 8, + connector_max_positions: list[int] | None = None, + connector_ff_bias: bool = True, + feature_extractor: LTX25FeatureExtractorV2 | None = None, + connectors: LTX25EmbeddingsConnectors | None = None, + ): + super().__init__() + self.feature_extractor = ( + feature_extractor + if feature_extractor is not None + else LTX25FeatureExtractorV2( + embedding_dim=embedding_dim, + num_layers=num_layers, + video_out_features=video_attention_heads * video_attention_head_dim, + audio_out_features=audio_attention_heads * audio_attention_head_dim, + ) + ) + self.connectors = ( + connectors + if connectors is not None + else LTX25EmbeddingsConnectors( + video_attention_heads=video_attention_heads, + video_attention_head_dim=video_attention_head_dim, + audio_attention_heads=audio_attention_heads, + audio_attention_head_dim=audio_attention_head_dim, + num_connector_layers=num_connector_layers, + connector_max_positions=connector_max_positions, + connector_ff_bias=connector_ff_bias, + ) + ) + + @property + def video_connector(self): + return self.connectors.video_connector + + @property + def audio_connector(self): + return self.connectors.audio_connector + + def create_embeddings( + self, + video_features: torch.Tensor, + audio_features: torch.Tensor | None, + additive_attention_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if audio_features is None: + raise ValueError("LTX-2.5 requires audio features for the audio connector.") + sort_indices, connector_mask = _right_pad_order(additive_attention_mask) + video_features = torch.gather(video_features, 1, sort_indices.unsqueeze(-1).expand_as(video_features)) + video_encoded, video_mask = self.video_connector(video_features, connector_mask) + binary_mask = (video_mask < 0.000001).to(torch.int64).reshape(video_encoded.shape[0], video_encoded.shape[1], 1) + video_encoded = video_encoded * binary_mask + audio_features = torch.gather(audio_features, 1, sort_indices.unsqueeze(-1).expand_as(audio_features)) + audio_encoded, _ = self.audio_connector(audio_features, connector_mask) + return video_encoded, audio_encoded, binary_mask.squeeze(-1) + + def process_hidden_states( + self, + hidden_states: tuple[torch.Tensor, ...], + attention_mask: torch.Tensor, + padding_side: str = "left", + ) -> LTX25EmbeddingsProcessorOutput: + video_features, audio_features = self.feature_extractor(hidden_states, attention_mask, padding_side) + additive_attention_mask = _convert_to_additive_mask(attention_mask, video_features.dtype) + video_encoding, audio_encoding, binary_mask = self.create_embeddings( + video_features, + audio_features, + additive_attention_mask, + ) + return LTX25EmbeddingsProcessorOutput(video_encoding, audio_encoding, binary_mask) diff --git a/diffsynth/models/ltx2_audio_vae.py b/diffsynth/models/ltx2_audio_vae.py index 8a58f9724..0c48583a3 100644 --- a/diffsynth/models/ltx2_audio_vae.py +++ b/diffsynth/models/ltx2_audio_vae.py @@ -821,14 +821,16 @@ class PerChannelStatistics(nn.Module): 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)) + # Frozen parameters, not buffers: disk-offload reload restores only named_parameters(), + # so persistent buffers would be skipped and strict load_state_dict would report them missing. + self.register_parameter("std-of-means", nn.Parameter(torch.empty(latent_channels), requires_grad=False)) + self.register_parameter("mean-of-means", nn.Parameter(torch.empty(latent_channels), requires_grad=False)) 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) + return (x * getattr(self, "std-of-means").to(x)) + getattr(self, "mean-of-means").to(x) def normalize(self, x: torch.Tensor) -> torch.Tensor: - return (x - self.get_buffer("mean-of-means").to(x)) / self.get_buffer("std-of-means").to(x) + return (x - getattr(self, "mean-of-means").to(x)) / getattr(self, "std-of-means").to(x) LATENT_DOWNSAMPLE_FACTOR = 4 @@ -1342,7 +1344,8 @@ def __init__( self.stride = stride self.padding = padding self.padding_mode = padding_mode - self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size)) + # Parameter, not buffer, so disk-offload reload (named_parameters() only) restores it from the checkpoint. + self.register_parameter("filter", nn.Parameter(kaiser_sinc_filter1d(cutoff, half_width, kernel_size), requires_grad=False)) def forward(self, x: torch.Tensor) -> torch.Tensor: _, n_channels, _ = x.shape @@ -1388,7 +1391,12 @@ def __init__( kernel_size=self.kernel_size, ) - self.register_buffer("filter", sinc_filter, persistent=persistent) + if persistent: + # Checkpoint-stored: parameter so disk-offload reload (named_parameters() only) restores it. + self.register_parameter("filter", nn.Parameter(sinc_filter, requires_grad=False)) + else: + # Not stored in the checkpoint: keep as a non-persistent buffer (recomputed above). + self.register_buffer("filter", sinc_filter, persistent=False) def forward(self, x: torch.Tensor) -> torch.Tensor: _, n_channels, _ = x.shape @@ -1701,8 +1709,10 @@ def __init__(self, filter_length: int, hop_length: int, win_length: int) -> None 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)) + # Parameters, not buffers: disk-offload reload restores only named_parameters(); + # the zeros here are overwritten by load_state_dict from the checkpoint. + self.register_parameter("forward_basis", nn.Parameter(torch.zeros(n_freqs * 2, 1, filter_length), requires_grad=False)) + self.register_parameter("inverse_basis", nn.Parameter(torch.zeros(n_freqs * 2, 1, filter_length), requires_grad=False)) def forward(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Compute magnitude and phase spectrogram from a batch of waveforms. @@ -1746,8 +1756,9 @@ def __init__( # Initialized to zeros; load_state_dict overwrites with the checkpoint's # exact bfloat16 filterbank (vocoder.mel_stft.mel_basis, shape [n_mels, n_freqs]). + # Parameter, not buffer, so disk-offload reload (named_parameters() only) restores it. n_freqs = filter_length // 2 + 1 - self.register_buffer("mel_basis", torch.zeros(n_mel_channels, n_freqs)) + self.register_parameter("mel_basis", nn.Parameter(torch.zeros(n_mel_channels, n_freqs), requires_grad=False)) def mel_spectrogram(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Compute log-mel spectrogram and auxiliary spectral quantities. diff --git a/diffsynth/models/ltx2_common.py b/diffsynth/models/ltx2_common.py index a658ec6ec..21270b884 100644 --- a/diffsynth/models/ltx2_common.py +++ b/diffsynth/models/ltx2_common.py @@ -266,6 +266,8 @@ class Modality: attention. ``None`` means unrestricted (full) attention between all tokens. Built incrementally by conditioning items; see :class:`~ltx_core.conditioning.types.attention_strength_wrapper.ConditioningItemAttentionStrengthWrapper`. + keyframes_mask: Optional per-token marker of shape ``(B, T, 1)``. Non-zero + entries receive the LTX 2.5 keyframe absolute-position embedding. """ latent: ( @@ -280,6 +282,7 @@ class Modality: enabled: bool = True context_mask: torch.Tensor | None = None attention_mask: torch.Tensor | None = None + keyframes_mask: torch.Tensor | None = None def to_denoised( diff --git a/diffsynth/models/ltx2_dit.py b/diffsynth/models/ltx2_dit.py index 9df0ed3a7..dbcc919a3 100644 --- a/diffsynth/models/ltx2_dit.py +++ b/diffsynth/models/ltx2_dit.py @@ -578,6 +578,23 @@ def forward(self, caption: torch.Tensor) -> torch.Tensor: hidden_states = self.linear_2(hidden_states) return hidden_states +KeyframesEmbeddingProvider = Callable[[], torch.Tensor | None] + + +def apply_keyframes_absolute_embedding( + hidden_states: torch.Tensor, + keyframes_mask: torch.Tensor | None, + embedding_provider: KeyframesEmbeddingProvider | None, +) -> torch.Tensor: + if embedding_provider is None or keyframes_mask is None: + return hidden_states + embedding = embedding_provider() + if embedding is None: + return hidden_states + mask = (keyframes_mask > 0).to(dtype=hidden_states.dtype) + return hidden_states + mask * embedding.to(device=hidden_states.device, dtype=hidden_states.dtype) + + @dataclass(frozen=True) class TransformerArgs: x: torch.Tensor @@ -611,6 +628,7 @@ def __init__( # noqa: PLR0913 rope_type: LTXRopeType, caption_projection: torch.nn.Module | None = None, prompt_adaln: AdaLayerNormSingle | None = None, + keyframes_embedding_provider: KeyframesEmbeddingProvider | None = None, ) -> None: self.patchify_proj = patchify_proj self.adaln = adaln @@ -624,6 +642,7 @@ def __init__( # noqa: PLR0913 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 @@ -717,6 +736,7 @@ def prepare( cross_modality: Modality | None = None, # noqa: ARG002 ) -> TransformerArgs: x = self.patchify_proj(modality.latent) + x = apply_keyframes_absolute_embedding(x, modality.keyframes_mask, self.keyframes_embedding_provider) batch_size = x.shape[0] timestep, embedded_timestep = self._prepare_timestep( modality.timesteps, self.adaln, batch_size, modality.latent.dtype @@ -773,6 +793,8 @@ def __init__( # noqa: PLR0913 av_ca_timestep_scale_multiplier: int, caption_projection: torch.nn.Module | None = None, prompt_adaln: AdaLayerNormSingle | None = None, + keyframes_embedding_provider: KeyframesEmbeddingProvider | None = None, + use_tokenwise_av_ca_scale_shift: bool = False, ) -> None: self.simple_preprocessor = TransformerArgsPreprocessor( patchify_proj=patchify_proj, @@ -787,12 +809,14 @@ def __init__( # noqa: PLR0913 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 + self.use_tokenwise_av_ca_scale_shift = use_tokenwise_av_ca_scale_shift def prepare( self, @@ -823,10 +847,11 @@ def prepare( ) cross_scale_shift_timestep, cross_gate_timestep = self._prepare_cross_attention_timestep( - timestep=cross_timestep, + timestep=modality.timesteps if self.use_tokenwise_av_ca_scale_shift else cross_timestep, timestep_scale_multiplier=self.simple_preprocessor.timestep_scale_multiplier, batch_size=transformer_args.x.shape[0], hidden_dtype=modality.latent.dtype, + gate_timestep=cross_timestep if self.use_tokenwise_av_ca_scale_shift else None, ) return replace( @@ -842,11 +867,14 @@ def _prepare_cross_attention_timestep( timestep_scale_multiplier: int, batch_size: int, hidden_dtype: torch.dtype, + gate_timestep: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Prepare cross attention timestep embeddings.""" timestep = timestep * timestep_scale_multiplier av_ca_factor = self.av_ca_timestep_scale_multiplier / timestep_scale_multiplier + # LTX-2.5 drives the scale/shift AdaLN per token while the gate stays scalar. + gate_input = (gate_timestep * timestep_scale_multiplier).flatten() if gate_timestep is not None else timestep.flatten() scale_shift_timestep, _ = self.cross_scale_shift_adaln( timestep.flatten(), @@ -854,7 +882,7 @@ def _prepare_cross_attention_timestep( ) scale_shift_timestep = scale_shift_timestep.view(batch_size, -1, scale_shift_timestep.shape[-1]) gate_noise_timestep, _ = self.cross_gate_adaln( - timestep.flatten() * av_ca_factor, + gate_input * av_ca_factor, hidden_dtype=hidden_dtype, ) gate_noise_timestep = gate_noise_timestep.view(batch_size, -1, gate_noise_timestep.shape[-1]) @@ -870,6 +898,7 @@ class TransformerConfig: context_dim: int apply_gated_attention: bool = False cross_attention_adaln: bool = False + ff_bias: bool = True class BasicAVTransformerBlock(torch.nn.Module): @@ -903,7 +932,7 @@ def __init__( norm_eps=norm_eps, apply_gated_attention=video.apply_gated_attention, ) - self.ff = FeedForward(video.dim, dim_out=video.dim) + 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)) @@ -926,7 +955,7 @@ def __init__( norm_eps=norm_eps, apply_gated_attention=audio.apply_gated_attention, ) - self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim) + 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)) @@ -1243,22 +1272,21 @@ def apply_cross_attention_adaln( class GELUApprox(torch.nn.Module): - def __init__(self, dim_in: int, dim_out: int) -> None: + def __init__(self, dim_in: int, dim_out: int, bias: bool = True) -> None: super().__init__() - self.proj = torch.nn.Linear(dim_in, dim_out) + 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") class FeedForward(torch.nn.Module): - def __init__(self, dim: int, dim_out: int, mult: int = 4) -> None: + 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) - - self.net = torch.nn.Sequential(project_in, torch.nn.Identity(), torch.nn.Linear(inner_dim, dim_out)) + 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) @@ -1309,6 +1337,11 @@ def __init__( # noqa: PLR0913 double_precision_rope: bool = True, apply_gated_attention: bool = False, cross_attention_adaln: bool = False, + use_prompt_adaln_single: bool = True, + ff_bias: bool = True, + audio_ff_bias: bool = True, + use_keyframes_abs_pos_embedding: bool = False, + use_tokenwise_av_ca_scale_shift: bool = False, ): super().__init__() self._enable_gradient_checkpointing = False @@ -1319,7 +1352,9 @@ def __init__( # noqa: PLR0913 self.positional_embedding_theta = positional_embedding_theta self.model_type = model_type self.cross_attention_adaln = cross_attention_adaln - cross_pe_max_pos = None + self.use_prompt_adaln_single = use_prompt_adaln_single + self.use_keyframes_abs_pos_embedding = use_keyframes_abs_pos_embedding + self.use_tokenwise_av_ca_scale_shift = use_tokenwise_av_ca_scale_shift if model_type.is_video_enabled(): if positional_embedding_max_pos is None: positional_embedding_max_pos = [20, 2048, 2048] @@ -1347,12 +1382,10 @@ def __init__( # noqa: PLR0913 ) 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, @@ -1362,12 +1395,17 @@ def __init__( # noqa: PLR0913 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, ) @property def _adaln_embedding_coefficient(self) -> int: return adaln_embedding_coefficient(self.cross_attention_adaln) + def _keyframes_embedding(self) -> torch.Tensor | None: + return getattr(self, "keyframes_abs_pos_embedding", None) + def _init_video( self, in_channels: int, @@ -1379,7 +1417,8 @@ def _init_video( # Video input components self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True) 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 + self.prompt_adaln_single = AdaLayerNormSingle(self.inner_dim, embedding_coefficient=2) if self.cross_attention_adaln and self.use_prompt_adaln_single else None + self.keyframes_abs_pos_embedding = torch.nn.Parameter(torch.zeros(1, self.inner_dim)) if self.use_keyframes_abs_pos_embedding else None # Video caption projection if caption_channels is not None: @@ -1406,7 +1445,7 @@ def _init_audio( self.audio_patchify_proj = torch.nn.Linear(in_channels, self.audio_inner_dim, bias=True) 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 + self.audio_prompt_adaln_single = AdaLayerNormSingle(self.audio_inner_dim, embedding_coefficient=2) if self.cross_attention_adaln and self.use_prompt_adaln_single else None # Audio caption projection if caption_channels is not None: @@ -1470,6 +1509,8 @@ def _init_preprocessors( 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, + use_tokenwise_av_ca_scale_shift=self.use_tokenwise_av_ca_scale_shift, ) self.audio_args_preprocessor = MultiModalTransformerArgsPreprocessor( patchify_proj=self.audio_patchify_proj, @@ -1489,6 +1530,7 @@ def _init_preprocessors( 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), + use_tokenwise_av_ca_scale_shift=self.use_tokenwise_av_ca_scale_shift, ) elif self.model_type.is_video_enabled(): self.video_args_preprocessor = TransformerArgsPreprocessor( @@ -1504,6 +1546,7 @@ def _init_preprocessors( 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( @@ -1530,6 +1573,8 @@ def _init_transformer_blocks( 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 = ( @@ -1540,6 +1585,7 @@ def _init_transformer_blocks( 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 @@ -1552,6 +1598,7 @@ def _init_transformer_blocks( 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 @@ -1672,12 +1719,26 @@ def _forward( ) return vx, ax - def forward(self, video_latents, video_positions, video_context, video_timesteps, audio_latents, audio_positions, audio_context, audio_timesteps, sigma, use_gradient_checkpointing=False, use_gradient_checkpointing_offload=False): + def forward( + self, + video_latents, + video_positions, + video_context, + video_timesteps, + audio_latents, + audio_positions, + audio_context, + audio_timesteps, + sigma, + use_gradient_checkpointing=False, + use_gradient_checkpointing_offload=False, + video_keyframes_mask=None, + ): cross_pe_max_pos = None if self.model_type.is_video_enabled() and self.model_type.is_audio_enabled(): cross_pe_max_pos = max(self.positional_embedding_max_pos[0], self.audio_positional_embedding_max_pos[0]) self._init_preprocessors(cross_pe_max_pos) - video = Modality(video_latents, sigma, video_timesteps, video_positions, video_context) + video = Modality(video_latents, sigma, video_timesteps, video_positions, video_context, keyframes_mask=video_keyframes_mask) if video_latents is not None else None audio = Modality(audio_latents, sigma, audio_timesteps, audio_positions, audio_context) if audio_latents is not None else None vx, ax = self._forward(video=video, audio=audio, perturbations=None, use_gradient_checkpointing=use_gradient_checkpointing, use_gradient_checkpointing_offload=use_gradient_checkpointing_offload) return vx, ax diff --git a/diffsynth/models/ltx2_video_vae.py b/diffsynth/models/ltx2_video_vae.py index a70dc80e8..026c63ff1 100644 --- a/diffsynth/models/ltx2_video_vae.py +++ b/diffsynth/models/ltx2_video_vae.py @@ -553,15 +553,17 @@ class PerChannelStatistics(nn.Module): 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)) + # Frozen parameters, not buffers: disk-offload reload restores only named_parameters(), + # so persistent buffers would be skipped and strict load_state_dict would report them missing. + self.register_parameter("std-of-means", nn.Parameter(torch.empty(latent_channels), requires_grad=False)) + self.register_parameter("mean-of-means", nn.Parameter(torch.empty(latent_channels), requires_grad=False)) 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( + return (x * getattr(self, "std-of-means").view(1, -1, 1, 1, 1).to(x)) + getattr(self, "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( + return (x - getattr(self, "mean-of-means").view(1, -1, 1, 1, 1).to(x)) / getattr(self, "std-of-means").view( 1, -1, 1, 1, 1).to(x) @@ -2187,6 +2189,7 @@ def decode( tile_overlap_in_pixels: Optional[int] = 128, tile_size_in_frames: Optional[int] = 128, tile_overlap_in_frames: Optional[int] = 24, + **kwargs, ) -> torch.Tensor: if tiled: tiling_config = TilingConfig( diff --git a/diffsynth/pipelines/ltx2_audio_video.py b/diffsynth/pipelines/ltx2_audio_video.py index 9cde576a3..61f24dc22 100644 --- a/diffsynth/pipelines/ltx2_audio_video.py +++ b/diffsynth/pipelines/ltx2_audio_video.py @@ -8,19 +8,19 @@ from PIL import Image from tqdm import tqdm from typing import Optional -from transformers import AutoImageProcessor, Gemma3Processor from ..core.device.npu_compatible_device import get_device_type -from ..diffusion import FlowMatchScheduler +from ..diffusion import AncestralFlowMatchScheduler, FlowMatchScheduler from ..core import ModelConfig from ..diffusion.base_pipeline import BasePipeline, PipelineUnit from ..models.ltx2_text_encoder import LTX2TextEncoder, LTX2TextEncoderPostModules, LTXVGemmaTokenizer from ..models.ltx2_dit import LTXModel -from ..models.ltx2_video_vae import LTX2VideoEncoder, LTX2VideoDecoder, VideoLatentPatchifier +from ..models.ltx2_video_vae import LTX2VideoEncoder, VideoLatentPatchifier from ..models.ltx2_audio_vae import LTX2AudioEncoder, LTX2AudioDecoder, LTX2Vocoder, AudioPatchifier, AudioProcessor from ..models.ltx2_upsampler import LTX2LatentUpsampler from ..models.ltx2_common import VideoLatentShape, AudioLatentShape, VideoPixelShape, get_pixel_coords, VIDEO_SCALE_FACTORS +from ..models.ltx25_text_encoder import LTX25GemmaTokenizer from ..utils.data.media_io_ltx2 import ltx2_preprocess from ..utils.data.audio import convert_to_stereo @@ -39,15 +39,16 @@ def __init__(self, device=get_device_type(), torch_dtype=torch.bfloat16): self.scheduler = FlowMatchScheduler("LTX-2") self.text_encoder: LTX2TextEncoder = None self.tokenizer: LTXVGemmaTokenizer = None - self.processor: Gemma3Processor = None self.text_encoder_post_modules: LTX2TextEncoderPostModules = None self.dit: LTXModel = None self.video_vae_encoder: LTX2VideoEncoder = None - self.video_vae_decoder: LTX2VideoDecoder = None + self.video_vae_decoder = None self.audio_vae_encoder: LTX2AudioEncoder = None self.audio_vae_decoder: LTX2AudioDecoder = None self.audio_vocoder: LTX2Vocoder = None self.upsampler: LTX2LatentUpsampler = None + self.duration_head = None + self.is_ltx25 = False self.video_patchifier: VideoLatentPatchifier = VideoLatentPatchifier(patch_size=1) self.audio_patchifier: AudioPatchifier = AudioPatchifier(patch_size=1) @@ -56,8 +57,9 @@ def __init__(self, device=get_device_type(), torch_dtype=torch.bfloat16): self.in_iteration_models = ("dit",) self.units = [ LTX2AudioVideoUnit_PipelineChecker(), - LTX2AudioVideoUnit_ShapeChecker(), LTX2AudioVideoUnit_PromptEmbedder(), + LTX2AudioVideoUnit_AutoDuration(), + LTX2AudioVideoUnit_ShapeChecker(), LTX2AudioVideoUnit_NoiseInitializer(), LTX2AudioVideoUnit_VideoRetakeEmbedder(), LTX2AudioVideoUnit_AudioRetakeEmbedder(), @@ -105,6 +107,20 @@ def __init__(self, device=get_device_type(), torch_dtype=torch.bfloat16): "pauses, incorrect timing, unnatural transitions, inconsistent framing, tilted camera, flat lighting, " "inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts." ), + "LTX-2.5": ( + "has_subtitles, has_blurbox, transition from black, transition to black, speech_ending_short, " + "blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, excessive noise, " + "grainy texture, poor lighting, flickering, motion blur, distorted proportions, unnatural skin tones, " + "deformed facial features, asymmetrical face, missing facial features, extra limbs, disfigured hands, " + "wrong hand count, artifacts around text, inconsistent perspective, camera shake, incorrect depth of " + "field, background too sharp, background clutter, distracting reflections, harsh shadows, inconsistent " + "lighting direction, color banding, cartoonish rendering, 3D CGI look, unrealistic materials, uncanny " + "valley effect, incorrect ethnicity, wrong gender, exaggerated expressions, wrong gaze direction, " + "mismatched lip sync, silent or muted audio, distorted voice, robotic voice, echo, background noise, " + "off-sync audio, incorrect dialogue, added dialogue, repetitive speech, jittery movement, awkward " + "pauses, incorrect timing, unnatural transitions, inconsistent framing, tilted camera, flat lighting, " + "inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts." + ), } @staticmethod @@ -123,19 +139,21 @@ def from_pretrained( # Fetch models pipe.text_encoder = model_pool.fetch_model("ltx2_text_encoder") + pipe.dit = model_pool.fetch_model("ltx2_dit") + pipe.is_ltx25 = getattr(pipe.dit, "use_tokenwise_av_ca_scale_shift", False) tokenizer_config.download_if_necessary() - pipe.tokenizer = LTXVGemmaTokenizer(tokenizer_path=tokenizer_config.path) - image_processor = AutoImageProcessor.from_pretrained(tokenizer_config.path, local_files_only=True) - pipe.processor = Gemma3Processor(image_processor=image_processor, tokenizer=pipe.tokenizer.tokenizer) - + tokenizer_class = LTX25GemmaTokenizer if pipe.is_ltx25 else LTXVGemmaTokenizer + pipe.tokenizer = tokenizer_class(tokenizer_path=tokenizer_config.path) pipe.text_encoder_post_modules = model_pool.fetch_model("ltx2_text_encoder_post_modules") - pipe.dit = model_pool.fetch_model("ltx2_dit") pipe.video_vae_encoder = model_pool.fetch_model("ltx2_video_vae_encoder") pipe.video_vae_decoder = model_pool.fetch_model("ltx2_video_vae_decoder") + if pipe.video_vae_decoder is None: + pipe.video_vae_decoder = model_pool.fetch_model("ltx25_diffusion_video_vae_decoder") pipe.audio_vae_decoder = model_pool.fetch_model("ltx2_audio_vae_decoder") pipe.audio_vocoder = model_pool.fetch_model("ltx2_audio_vocoder") pipe.upsampler = model_pool.fetch_model("ltx2_latent_upsampler") pipe.audio_vae_encoder = model_pool.fetch_model("ltx2_audio_vae_encoder") + pipe.duration_head = model_pool.fetch_model("ltx25_duration_head") # Stage 2 if stage2_lora_config is not None: @@ -159,8 +177,9 @@ def denoise_stage(self, inputs_shared, inputs_posi, inputs_nega, units, cfg_scal self.model_fn, cfg_scale, inputs_shared, inputs_posi, inputs_nega, **models, timestep=timestep, progress_id=progress_id ) - inputs_shared["video_latents"] = self.step(self.scheduler, inputs_shared["video_latents"], progress_id=progress_id, noise_pred=noise_pred_video, - inpaint_mask=inputs_shared.get("denoise_mask_video", None), input_latents=inputs_shared.get("input_latents_video", None), **inputs_shared) + if inputs_shared.get("video_latents") is not None: + inputs_shared["video_latents"] = self.step(self.scheduler, inputs_shared["video_latents"], progress_id=progress_id, noise_pred=noise_pred_video, + inpaint_mask=inputs_shared.get("denoise_mask_video", None), input_latents=inputs_shared.get("input_latents_video", None), **inputs_shared) inputs_shared["audio_latents"] = self.step(self.scheduler, inputs_shared["audio_latents"], progress_id=progress_id, noise_pred=noise_pred_audio, inpaint_mask=inputs_shared.get("denoise_mask_audio", None), input_latents=inputs_shared.get("input_latents_audio", None), **inputs_shared) return inputs_shared, inputs_posi, inputs_nega @@ -194,6 +213,10 @@ def __call__( width: int = 768, num_frames: int = 121, frame_rate: int = 24, + auto_duration: bool = False, + auto_duration_min_seconds: float = 1.0, + auto_duration_max_seconds: float = 20.0, + audio_only: bool = False, # Classifier-free guidance cfg_scale: float = 3.0, # Scheduler @@ -213,14 +236,12 @@ def __call__( progress_bar_cmd=tqdm, ): # Scheduler + if self.is_ltx25 and use_distilled_pipeline: + self.scheduler = AncestralFlowMatchScheduler(noise_seed=seed + 10000, rand_device=rand_device) self.scheduler.set_timesteps(num_inference_steps, denoising_strength=denoising_strength, special_case="distilled_stage1" if use_distilled_pipeline else None) # Inputs - inputs_posi = { - "prompt": prompt, - } - inputs_nega = { - "negative_prompt": negative_prompt, - } + inputs_posi = {"prompt": prompt} + inputs_nega = {"negative_prompt": negative_prompt} inputs_shared = { "input_images": input_images, "input_images_indexes": input_images_indexes, "input_images_strength": input_images_strength, "retake_video": retake_video, "retake_video_regions": retake_video_regions, @@ -228,6 +249,8 @@ def __call__( "in_context_videos": in_context_videos, "in_context_downsample_factor": in_context_downsample_factor, "seed": seed, "rand_device": rand_device, "height": height, "width": width, "num_frames": num_frames, "frame_rate": frame_rate, + "auto_duration": auto_duration, "auto_duration_min_seconds": auto_duration_min_seconds, "auto_duration_max_seconds": auto_duration_max_seconds, + "audio_only": audio_only, "cfg_scale": cfg_scale, "tiled": tiled, "tile_size_in_pixels": tile_size_in_pixels, "tile_overlap_in_pixels": tile_overlap_in_pixels, "tile_size_in_frames": tile_size_in_frames, "tile_overlap_in_frames": tile_overlap_in_frames, @@ -239,10 +262,14 @@ def __call__( # Stage 2 inputs_shared, inputs_posi, inputs_nega = self.denoise_stage(inputs_shared, inputs_posi, inputs_nega, self.stage2_units, 1.0, progress_bar_cmd, not inputs_shared["use_two_stage_pipeline"]) # Decode - self.load_models_to_device(['video_vae_decoder']) - video = self.video_vae_decoder.decode(inputs_shared["video_latents"], tiled, tile_size_in_pixels, tile_overlap_in_pixels, tile_size_in_frames, tile_overlap_in_frames) - video = self.vae_output_to_video(video) - self.load_models_to_device(['audio_vae_decoder', 'audio_vocoder']) + video = None + if not inputs_shared.get("audio_only", False): + self.load_models_to_device(["video_vae_decoder"]) + video = self.video_vae_decoder.decode( + inputs_shared["video_latents"], tiled=tiled, tile_size_in_pixels=tile_size_in_pixels, tile_overlap_in_pixels=tile_overlap_in_pixels, + tile_size_in_frames=tile_size_in_frames, tile_overlap_in_frames=tile_overlap_in_frames, seed=seed, rand_device=rand_device) + video = self.vae_output_to_video(video) + self.load_models_to_device(["audio_vae_decoder", "audio_vocoder"]) decoded_audio = self.audio_vae_decoder(inputs_shared["audio_latents"]) decoded_audio = self.audio_vocoder(decoded_audio) decoded_audio = self.output_audio_format_check(decoded_audio) @@ -271,6 +298,30 @@ def process(self, pipe: LTX2AudioVideoPipeline, inputs_shared, inputs_posi, inpu return inputs_shared, inputs_posi, inputs_nega +class LTX2AudioVideoUnit_AutoDuration(PipelineUnit): + def __init__(self): + super().__init__(take_over=True, onload_model_names=("duration_head",)) + + @staticmethod + def seconds_to_num_frames(seconds, frame_rate, min_seconds, max_seconds): + min_frames = round(min_seconds * frame_rate) + max_frames = round(max_seconds * frame_rate) + raw_frames = max(min_frames, min(round(seconds * frame_rate), max_frames)) + frames = ((raw_frames - 1) // 8) * 8 + 1 + if frames < min_frames: + frames = min(-(-(min_frames - 1) // 8) * 8 + 1, max_frames) + return frames + + def process(self, pipe: LTX2AudioVideoPipeline, inputs_shared, inputs_posi, inputs_nega): + if not inputs_shared.get("auto_duration", False) or pipe.duration_head is None: + return inputs_shared, inputs_posi, inputs_nega + min_seconds, max_seconds = sorted(min(max(seconds, 1e-6), 20.0) for seconds in (inputs_shared["auto_duration_min_seconds"], inputs_shared["auto_duration_max_seconds"])) + pipe.load_models_to_device(self.onload_model_names) + seconds = float(pipe.duration_head(inputs_posi["video_context"], inputs_posi["audio_context"]).item()) + inputs_shared["num_frames"] = self.seconds_to_num_frames(seconds, inputs_shared["frame_rate"], min_seconds, max_seconds) + return inputs_shared, inputs_posi, inputs_nega + + class LTX2AudioVideoUnit_ShapeChecker(PipelineUnit): """ For two-stage pipelines, the resolution must be divisible by 64. @@ -329,19 +380,30 @@ def process(self, pipe: LTX2AudioVideoPipeline, prompt: str): class LTX2AudioVideoUnit_NoiseInitializer(PipelineUnit): def __init__(self): super().__init__( - input_params=("height", "width", "num_frames", "seed", "rand_device", "frame_rate"), - output_params=("video_noise", "audio_noise", "video_positions", "audio_positions", "video_latent_shape", "audio_latent_shape") + input_params=("height", "width", "num_frames", "seed", "rand_device", "frame_rate", "audio_only"), + output_params=("video_noise", "audio_noise", "video_positions", "audio_positions", "video_latent_shape", "audio_latent_shape", "video_keyframes_mask") ) - def process_stage(self, pipe: LTX2AudioVideoPipeline, height, width, num_frames, seed, rand_device, frame_rate=24.0): + def build_video_keyframes_mask(self, pipe, video_latent_shape): + if not pipe.is_ltx25: + return None + video_keyframes_mask = torch.zeros(video_latent_shape.batch, 1, video_latent_shape.frames, video_latent_shape.height, video_latent_shape.width, dtype=torch.float32, device=pipe.device) + video_keyframes_mask[:, :, 0] = 1.0 + return video_keyframes_mask + + def process_stage(self, pipe: LTX2AudioVideoPipeline, height, width, num_frames, seed, rand_device, frame_rate=24.0, audio_only=False): video_pixel_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate) video_latent_shape = VideoLatentShape.from_pixel_shape(shape=video_pixel_shape, latent_channels=128) - video_noise = pipe.generate_noise(video_latent_shape.to_torch_shape(), seed=seed, rand_device=rand_device) - - latent_coords = pipe.video_patchifier.get_patch_grid_bounds(output_shape=video_latent_shape, device=pipe.device) - video_positions = get_pixel_coords(latent_coords, VIDEO_SCALE_FACTORS, True).float() - video_positions[:, 0, ...] = video_positions[:, 0, ...] / frame_rate - video_positions = video_positions.to(pipe.torch_dtype) + video_noise = video_positions = video_keyframes_mask = None + if not audio_only: + video_noise = pipe.generate_noise(video_latent_shape.to_torch_shape(), seed=seed, rand_device=rand_device) + + latent_coords = pipe.video_patchifier.get_patch_grid_bounds(output_shape=video_latent_shape, device=pipe.device) + video_positions = get_pixel_coords(latent_coords, VIDEO_SCALE_FACTORS, True).float() + video_positions[:, 0, ...] = video_positions[:, 0, ...] / frame_rate + if not pipe.is_ltx25: + video_positions = video_positions.to(pipe.torch_dtype) + video_keyframes_mask = self.build_video_keyframes_mask(pipe, video_latent_shape) audio_latent_shape = AudioLatentShape.from_video_pixel_shape(video_pixel_shape) audio_noise = pipe.generate_noise(audio_latent_shape.to_torch_shape(), seed=seed, rand_device=rand_device) @@ -352,11 +414,12 @@ def process_stage(self, pipe: LTX2AudioVideoPipeline, height, width, num_frames, "video_positions": video_positions, "audio_positions": audio_positions, "video_latent_shape": video_latent_shape, - "audio_latent_shape": audio_latent_shape + "audio_latent_shape": audio_latent_shape, + "video_keyframes_mask": video_keyframes_mask, } - def process(self, pipe: LTX2AudioVideoPipeline, height, width, num_frames, seed, rand_device, frame_rate=24.0): - return self.process_stage(pipe, height, width, num_frames, seed, rand_device, frame_rate) + def process(self, pipe: LTX2AudioVideoPipeline, height, width, num_frames, seed, rand_device, frame_rate=24.0, audio_only=False): + return self.process_stage(pipe, height, width, num_frames, seed, rand_device, frame_rate, audio_only) class LTX2AudioVideoUnit_InputVideoEmbedder(PipelineUnit): @@ -531,7 +594,8 @@ def process( latent_coords = pipe.video_patchifier.get_patch_grid_bounds(output_shape=VideoLatentShape.from_torch_shape(latents.shape), device=pipe.device) video_positions = get_pixel_coords(latent_coords, VIDEO_SCALE_FACTORS, False).float() video_positions[:, 0, ...] = (video_positions[:, 0, ...] + index) / frame_rate - video_positions = video_positions.to(pipe.torch_dtype) + if not pipe.is_ltx25: + video_positions = video_positions.to(pipe.torch_dtype) frame_conditions["ref_frames_latents"].append(latents) frame_conditions["ref_frames_positions"].append(video_positions) if len(frame_conditions["ref_frames_latents"]) == 0: @@ -578,7 +642,8 @@ def process(self, pipe: LTX2AudioVideoPipeline, in_context_videos, height, width video_positions[:, 0, ...] = video_positions[:, 0, ...] / frame_rate video_positions[:, 1, ...] *= in_context_downsample_factor # height axis video_positions[:, 2, ...] *= in_context_downsample_factor # width axis - video_positions = video_positions.to(pipe.torch_dtype) + if not pipe.is_ltx25: + video_positions = video_positions.to(pipe.torch_dtype) latents.append(in_context_latents) positions.append(video_positions) @@ -619,8 +684,10 @@ def __init__(self): ) def process(self, pipe: LTX2AudioVideoPipeline, video_latents, video_noise, audio_latents, audio_noise): + pipe.scheduler = FlowMatchScheduler("LTX-2") pipe.scheduler.set_timesteps(special_case="stage2") - video_latents = pipe.scheduler.add_noise(video_latents, video_noise, pipe.scheduler.timesteps[0]) + if video_latents is not None and video_noise is not None: + video_latents = pipe.scheduler.add_noise(video_latents, video_noise, pipe.scheduler.timesteps[0]) audio_latents = pipe.scheduler.add_noise(audio_latents, audio_noise, pipe.scheduler.timesteps[0]) return {"video_latents": video_latents, "audio_latents": audio_latents} @@ -630,7 +697,7 @@ def __init__(self): super().__init__( input_params=("video_latents",), output_params=("video_latents",), - onload_model_names=("upsampler",), + onload_model_names=("upsampler", "video_vae_encoder"), ) def process(self, pipe: LTX2AudioVideoPipeline, video_latents): @@ -667,6 +734,8 @@ def model_fn_ltx2( # Audio Inputs input_latents_audio=None, denoise_mask_audio=None, + # LTX-2.5 keyframe class embedding + video_keyframes_mask=None, # Gradient Checkpointing use_gradient_checkpointing=False, use_gradient_checkpointing_offload=False, @@ -674,11 +743,16 @@ def model_fn_ltx2( ): timestep = timestep.float() / 1000. - # patchify - b, c_v, f, h, w = video_latents.shape - video_latents = video_patchifier.patchify(video_latents) - seq_len_video = video_latents.shape[1] - video_timesteps = timestep.repeat(1, video_latents.shape[1], 1) + video_timesteps = None + if video_latents is not None: + # patchify + b, c_v, f, h, w = video_latents.shape + video_latents = video_patchifier.patchify(video_latents) + if video_keyframes_mask is not None: + video_latents = video_latents.transpose(1, 2).contiguous().transpose(1, 2) + video_keyframes_mask = video_patchifier.patchify(video_keyframes_mask) + seq_len_video = video_latents.shape[1] + video_timesteps = timestep.repeat(1, video_latents.shape[1], 1) # Frist frame conditioning by replacing the video latents if input_latents_video is not None: denoise_mask_video = video_patchifier.patchify(denoise_mask_video) @@ -697,6 +771,9 @@ def model_fn_ltx2( video_latents = torch.cat([video_latents, ref_frames_latent], dim=1) video_positions = torch.cat([video_positions, ref_frames_position], dim=2) video_timesteps = torch.cat([video_timesteps, ref_frames_timestep], dim=1) + if video_keyframes_mask is not None: + ref_keyframes_mask = torch.zeros(ref_frames_latent.shape[0], ref_frames_latent.shape[1], 1, dtype=video_keyframes_mask.dtype, device=video_keyframes_mask.device) + video_keyframes_mask = torch.cat([video_keyframes_mask, ref_keyframes_mask], dim=1) if audio_latents is not None: _, c_a, _, mel_bins = audio_latents.shape @@ -719,12 +796,11 @@ def model_fn_ltx2( audio_context=audio_context, audio_timesteps=audio_timesteps, sigma=timestep, + video_keyframes_mask=video_keyframes_mask, use_gradient_checkpointing=use_gradient_checkpointing, use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, ) - vx = vx[:, :seq_len_video, ...] - # unpatchify - vx = video_patchifier.unpatchify_video(vx, f, h, w) - ax = audio_patchifier.unpatchify_audio(ax, c_a, mel_bins) if ax is not None else None + vx = video_patchifier.unpatchify_video(vx[:, :seq_len_video, ...], f, h, w) if vx is not None else 0 + ax = audio_patchifier.unpatchify_audio(ax, c_a, mel_bins) if ax is not None else 0 return vx, ax diff --git a/diffsynth/utils/state_dict_converters/ltx25_diffusion_video_vae.py b/diffsynth/utils/state_dict_converters/ltx25_diffusion_video_vae.py new file mode 100644 index 000000000..e6d84247b --- /dev/null +++ b/diffsynth/utils/state_dict_converters/ltx25_diffusion_video_vae.py @@ -0,0 +1,21 @@ +def LTX25DiffusionVideoDecoderStateDictConverter(state_dict): + """Select DiffVAE weights and rename only legacy timestep-MLP keys. + + The flat decoder keeps fused ``qkv`` parameters, so this converter is safe for + lazy ``DiskMap`` inputs and never reads tensor metadata or values. + """ + converted = {} + for source_name in state_dict: + if source_name.startswith("decoder."): + name = source_name.removeprefix("decoder.") + elif source_name.startswith("per_channel_statistics."): + name = source_name + else: + continue + + if name.startswith("coarse_") or name.endswith((".gate_msa", ".gate_mlp", ".gate_ctx")): + continue + name = name.replace("t_embedder.mlp.0.", "t_embedder.timestep_embedder.linear_1.") + name = name.replace("t_embedder.mlp.2.", "t_embedder.timestep_embedder.linear_2.") + converted[name] = state_dict[source_name] + return converted diff --git a/diffsynth/utils/state_dict_converters/ltx25_duration_head.py b/diffsynth/utils/state_dict_converters/ltx25_duration_head.py new file mode 100644 index 000000000..f2a52ea49 --- /dev/null +++ b/diffsynth/utils/state_dict_converters/ltx25_duration_head.py @@ -0,0 +1,6 @@ +def LTX25DurationHeadStateDictConverter(state_dict): + return { + name.removeprefix("duration_head."): state_dict[name] + for name in state_dict + if name.startswith("duration_head.") + } diff --git a/diffsynth/utils/state_dict_converters/ltx25_text_encoder.py b/diffsynth/utils/state_dict_converters/ltx25_text_encoder.py new file mode 100644 index 000000000..e27f0f29d --- /dev/null +++ b/diffsynth/utils/state_dict_converters/ltx25_text_encoder.py @@ -0,0 +1,16 @@ +def LTX25TextEncoderStateDictConverter(state_dict): + state_dict_ = {} + for name in state_dict: + if name.startswith("model."): + new_name = "model.language_model." + name.removeprefix("model.") + elif name.startswith("vision_model."): + new_name = "model.embed_vision." + name.removeprefix("vision_model.") + elif name.startswith("multi_modal_projector."): + new_name = "model.embed_vision.multimodal_embedder." + name.removeprefix("multi_modal_projector.") + elif name.startswith("audio_projector."): + new_name = "model.embed_audio." + name.removeprefix("audio_projector.") + else: + continue + state_dict_[new_name] = state_dict[name] + state_dict_["lm_head.weight"] = state_dict_["model.language_model.embed_tokens.weight"] + return state_dict_ diff --git a/diffsynth/utils/state_dict_converters/ltx2_video_vae.py b/diffsynth/utils/state_dict_converters/ltx2_video_vae.py index 53df15e54..492bc49ce 100644 --- a/diffsynth/utils/state_dict_converters/ltx2_video_vae.py +++ b/diffsynth/utils/state_dict_converters/ltx2_video_vae.py @@ -4,7 +4,9 @@ def LTX2VideoEncoderStateDictConverter(state_dict): if name.startswith("vae.encoder."): new_name = name.replace("vae.encoder.", "") state_dict_[new_name] = state_dict[name] - elif name.startswith("vae.per_channel_statistics."): + elif name.startswith("encoder."): + state_dict_[name.removeprefix("encoder.")] = state_dict[name] + elif name.startswith("vae.per_channel_statistics.") or name.startswith("per_channel_statistics."): new_name = name.replace("vae.per_channel_statistics.", "per_channel_statistics.") if new_name not in ["per_channel_statistics.channel", "per_channel_statistics.mean-of-stds", "per_channel_statistics.mean-of-stds_over_std-of-means"]: state_dict_[new_name] = state_dict[name] @@ -17,7 +19,9 @@ def LTX2VideoDecoderStateDictConverter(state_dict): if name.startswith("vae.decoder."): new_name = name.replace("vae.decoder.", "") state_dict_[new_name] = state_dict[name] - elif name.startswith("vae.per_channel_statistics."): + elif name.startswith("decoder."): + state_dict_[name.removeprefix("decoder.")] = state_dict[name] + elif name.startswith("vae.per_channel_statistics.") or name.startswith("per_channel_statistics."): new_name = name.replace("vae.per_channel_statistics.", "per_channel_statistics.") if new_name not in ["per_channel_statistics.channel", "per_channel_statistics.mean-of-stds", "per_channel_statistics.mean-of-stds_over_std-of-means"]: state_dict_[new_name] = state_dict[name] diff --git a/docs/en/Model_Details/LTX-2.md b/docs/en/Model_Details/LTX-2.md index 6760cdd89..bda9a5b45 100644 --- a/docs/en/Model_Details/LTX-2.md +++ b/docs/en/Model_Details/LTX-2.md @@ -39,7 +39,7 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), @@ -85,6 +85,16 @@ write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24 |[Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up)||[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Jib-Up.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Jib-Up.py)|-|-|-|-| |[Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down)||[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Jib-Down.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Jib-Down.py)|-|-|-|-| |[Lightricks/LTX-2-19b-LoRA-Camera-Control-Static](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Static)||[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Static.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Static.py)|-|-|-|-| +|[Lightricks/LTX-2.5: OneStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|-|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py)|[code](/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py)|[code](/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py)| +|[Lightricks/LTX-2.5: TwoStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|-|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: OneStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: TwoStagePipeline-A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_audio`,`audio_sample_rate`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: TwoStagePipeline-Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_video`,`retake_video_regions`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py)|-|-|-|-| +|[Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`audio_only=True`|[code](/examples/ltx2/model_inference/LTX-2.5-T2A.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py)|-|-|-|-| +|[Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`auto_duration`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py)|-|-|-|-| +|[Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler)|`in_context_videos`,`in_context_downsample_factor`|[code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|-|-|-|-| +|[Lightricks/LTX-2.5: INT8-ConvRot](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|INT8 DiT + INT8 Gemma4|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py)|-|-|-|-| ## Model Inference @@ -94,16 +104,19 @@ Input parameters for `LTX2AudioVideoPipeline` inference include: * `prompt`: Prompt describing the content appearing in the video. * `negative_prompt`: Negative prompt describing content that should not appear in the video, default value is `""`. -* `cfg_scale`: Classifier-free guidance parameter, default value is 3.0. +* `denoising_strength`: Denoising strength, range is 0～1, default value is 1.0. * `input_images`: List of input images for image-to-video generation. * `input_images_indexes`: Frame index list of input images in the video. * `input_images_strength`: Strength of input images, default value is 1.0. -* `denoising_strength`: Denoising strength, range is 0～1, default value is 1.0. * `seed`: Random seed. Default is `None`, which means completely random. * `rand_device`: Computing device for generating random Gaussian noise matrix, default is `"cpu"`. When set to `cuda`, different results will be generated on different GPUs. * `height`: Video height, must be a multiple of 32 (single-stage) or 64 (two-stage). * `width`: Video width, must be a multiple of 32 (single-stage) or 64 (two-stage). * `num_frames`: Number of video frames, default value is 121, must be a multiple of 8 + 1. +* `auto_duration`: Predict the clip duration from the prompt. Defaults to `False`. When enabled, `num_frames` is not required and the Duration Head must be loaded. +* `auto_duration_min_seconds` / `auto_duration_max_seconds`: Lower and upper bounds (seconds) for the predicted duration. Default to 1.0 and 20.0. +* `audio_only`: Whether to generate audio only (T2A). Defaults to `False`; when `True`, the video VAE and latent upsampler are not required. +* `cfg_scale`: Classifier-free guidance parameter, default value is 3.0. * `num_inference_steps`: Number of inference steps, default value is 40. * `tiled`: Whether to enable VAE tiling inference, default is `True`. When set to `True`, it can significantly reduce VRAM usage during VAE encoding/decoding stages, with slight errors and minor inference time extension. * `tile_size_in_pixels`: Pixel tiling size during VAE encoding/decoding stages, default is 512. @@ -114,6 +127,7 @@ Input parameters for `LTX2AudioVideoPipeline` inference include: * `use_distilled_pipeline`: Whether to use distilled pipeline, default is `False`. * `progress_bar_cmd`: Progress bar, default is `tqdm.tqdm`. Can be set to `lambda x:x` to hide the progress bar. + If VRAM is insufficient, please enable [VRAM Management](../Pipeline_Usage/VRAM_management.md). We provide recommended low VRAM configurations for each model in the example code, see the table in the previous "Supported Inference Scripts" section. ## Model Training diff --git a/docs/en/index.rst b/docs/en/index.rst index 92678c525..0780fd166 100644 --- a/docs/en/index.rst +++ b/docs/en/index.rst @@ -33,6 +33,7 @@ Welcome to DiffSynth-Studio's Documentation Model_Details/Z-Image Model_Details/Anima Model_Details/LTX-2 + Model_Details/LTX-2.5 Model_Details/ERNIE-Image Model_Details/JoyAI-Image Model_Details/ACE-Step diff --git a/docs/zh/Model_Details/LTX-2.md b/docs/zh/Model_Details/LTX-2.md index 6ba595be5..0261f103a 100644 --- a/docs/zh/Model_Details/LTX-2.md +++ b/docs/zh/Model_Details/LTX-2.md @@ -39,7 +39,7 @@ pipe = LTX2AudioVideoPipeline.from_pretrained( model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), @@ -85,6 +85,16 @@ write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24 |[Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up)||[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Jib-Up.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Jib-Up.py)|-|-|-|-| |[Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down)||[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Jib-Down.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Jib-Down.py)|-|-|-|-| |[Lightricks/LTX-2-19b-LoRA-Camera-Control-Static](https://www.modelscope.cn/models/Lightricks/LTX-2-19b-LoRA-Camera-Control-Static)||[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference/LTX-2-T2AV-Camera-Control-Static.py)|[code](https://github.com/modelscope/DiffSynth-Studio/blob/main/examples/ltx2/model_inference_low_vram/LTX-2-T2AV-Camera-Control-Static.py)|-|-|-|-| +|[Lightricks/LTX-2.5: OneStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|-|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py)|[code](/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py)|[code](/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh)|[code](/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py)| +|[Lightricks/LTX-2.5: TwoStagePipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|-|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: OneStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: TwoStagePipeline-I2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`input_images`,`input_images_indexes`|[code](/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: TwoStagePipeline-A2V](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_audio`,`audio_sample_rate`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py)|-|-|-|-| +|[Lightricks/LTX-2.5: TwoStagePipeline-Retake](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`retake_video`,`retake_video_regions`,`stage2_lora_config`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py)|-|-|-|-| +|[Lightricks/LTX-2.5: T2A](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`audio_only=True`|[code](/examples/ltx2/model_inference/LTX-2.5-T2A.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py)|-|-|-|-| +|[Lightricks/LTX-2.5: DistilledPipeline-T2AV](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|`auto_duration`|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py)|-|-|-|-| +|[Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler](https://www.modelscope.cn/models/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler)|`in_context_videos`,`in_context_downsample_factor`|[code](/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py)|-|-|-|-| +|[Lightricks/LTX-2.5: INT8-ConvRot](https://www.modelscope.cn/models/Lightricks/LTX-2.5)|INT8 DiT + INT8 Gemma4|[code](/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py)|[code](/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py)|-|-|-|-| ## 模型推理 @@ -94,16 +104,19 @@ write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24 * `prompt`: 提示词，描述视频中出现的内容。 * `negative_prompt`: 负向提示词，描述视频中不应该出现的内容，默认值为 `""`。 -* `cfg_scale`: Classifier-free guidance 的参数，默认值为 3.0。 +* `denoising_strength`: 去噪强度，范围是 0～1，默认值为 1.0。 * `input_images`: 输入图像列表，用于图生视频。 * `input_images_indexes`: 输入图像在视频中的帧索引列表。 * `input_images_strength`: 输入图像的强度，默认值为 1.0。 -* `denoising_strength`: 去噪强度，范围是 0～1，默认值为 1.0。 * `seed`: 随机种子。默认为 `None`，即完全随机。 * `rand_device`: 生成随机高斯噪声矩阵的计算设备，默认为 `"cpu"`。当设置为 `cuda` 时，在不同 GPU 上会导致不同的生成结果。 * `height`: 视频高度，需保证高度为 32 的倍数（单阶段）或 64 的倍数（两阶段）。 * `width`: 视频宽度，需保证宽度为 32 的倍数（单阶段）或 64 的倍数（两阶段）。 * `num_frames`: 视频帧数，默认值为 121，需保证为 8 的倍数 + 1。 +* `auto_duration`: 是否根据提示词自动预测视频时长，默认为 `False`。开启后无需传入 `num_frames`，需要加载 Duration Head。 +* `auto_duration_min_seconds` / `auto_duration_max_seconds`: 自动时长的上下界（秒），默认为 1.0 和 20.0。 +* `audio_only`: 是否只生成音频（T2A），默认为 `False`。设置为 `True` 时无需加载视频 VAE 与 latent upsampler。 +* `cfg_scale`: Classifier-free guidance 的参数，默认值为 3.0。 * `num_inference_steps`: 推理次数，默认值为 40。 * `tiled`: 是否启用 VAE 分块推理，默认为 `True`。设置为 `True` 时可显著减少 VAE 编解码阶段的显存占用，会产生少许误差，以及少量推理时间延长。 * `tile_size_in_pixels`: VAE 编解码阶段的像素分块大小，默认为 512。 @@ -114,7 +127,7 @@ write_video_audio_ltx2(video=video, audio=audio, output_path='video.mp4', fps=24 * `use_distilled_pipeline`: 是否使用蒸馏管道，默认为 `False`。 * `progress_bar_cmd`: 进度条，默认为 `tqdm.tqdm`。可通过设置为 `lambda x:x` 来屏蔽进度条。 -如果显存不足，请开启[显存管理](../Pipeline_Usage/VRAM_management.md)，我们在示例代码中提供了每个模型推荐的低显存配置，详见前文"支持的推理脚本"中的表格。 +如果显存不足，请开启[显存管理](../Pipeline_Usage/VRAM_management.md)，我们在示例代码中提供了每个模型推荐的低显存配置，详见前文"模型总览"中的表格。 ## 模型训练 diff --git a/docs/zh/index.rst b/docs/zh/index.rst index 8a85d49e8..5201bc755 100644 --- a/docs/zh/index.rst +++ b/docs/zh/index.rst @@ -33,6 +33,7 @@ Model_Details/Z-Image Model_Details/Anima Model_Details/LTX-2 + Model_Details/LTX-2.5 Model_Details/ERNIE-Image Model_Details/JoyAI-Image Model_Details/ACE-Step diff --git a/examples/ltx2/model_inference/LTX-2.3-A2V-TwoStage.py b/examples/ltx2/model_inference/LTX-2.3-A2V-TwoStage.py index ad0b91834..e474c2f24 100644 --- a/examples/ltx2/model_inference/LTX-2.3-A2V-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.3-A2V-TwoStage.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference/LTX-2.3-I2AV-DistilledPipeline.py b/examples/ltx2/model_inference/LTX-2.3-I2AV-DistilledPipeline.py index d23460f5e..81274195c 100644 --- a/examples/ltx2/model_inference/LTX-2.3-I2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference/LTX-2.3-I2AV-DistilledPipeline.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), ) diff --git a/examples/ltx2/model_inference/LTX-2.3-I2AV-TwoStage.py b/examples/ltx2/model_inference/LTX-2.3-I2AV-TwoStage.py index 0dd0854c5..03c5c1dd5 100644 --- a/examples/ltx2/model_inference/LTX-2.3-I2AV-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.3-I2AV-TwoStage.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference/LTX-2.3-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference/LTX-2.3-T2AV-DistilledPipeline.py index 047a0b0d5..03e9c6d48 100644 --- a/examples/ltx2/model_inference/LTX-2.3-T2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference/LTX-2.3-T2AV-DistilledPipeline.py @@ -18,7 +18,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), ) diff --git a/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py b/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py index 5708f0d6f..9c923018d 100644 --- a/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py +++ b/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Union-Control.py b/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Union-Control.py index 25883f681..ecc38e757 100644 --- a/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Union-Control.py +++ b/examples/ltx2/model_inference/LTX-2.3-T2AV-IC-LoRA-Union-Control.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage-Retake.py b/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage-Retake.py index d241f69da..dc1667379 100644 --- a/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage-Retake.py +++ b/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage-Retake.py @@ -21,7 +21,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage.py b/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage.py index 52345894e..ca172e7f9 100644 --- a/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage.py +++ b/examples/ltx2/model_inference/LTX-2.3-T2AV-TwoStage.py @@ -18,7 +18,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py new file mode 100644 index 000000000..64b5c6e7b --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-A2V-TwoStage.py @@ -0,0 +1,61 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from diffsynth.utils.data.audio import read_audio +from modelscope import dataset_snapshot_download + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], + stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), + stage2_lora_strength=1.0, +) + +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree." +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames, frame_rate = 512 * 2, 768 * 2, 121, 24 +duration = num_frames / frame_rate +audio, audio_sample_rate = read_audio("data/example_video_dataset/ltx2/sing.MP3", start_time=1, duration=duration) +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + retake_audio=audio, + audio_sample_rate=audio_sample_rate, + seed=42, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + tiled=True, + tile_size_in_frames=80, + use_two_stage_pipeline=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_twostage_a2v.mp4", + fps=frame_rate, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py b/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py new file mode 100644 index 000000000..e342d0339 --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-I2AV-OneStage.py @@ -0,0 +1,55 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from PIL import Image +from modelscope import dataset_snapshot_download + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ], +) +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +prompt = "Two cute orange cats, wearing boxing gloves, stand in a boxing ring and fight each other. They are punching each other fast and yelling: 'I will win!'" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512, 768, 121 +first_frame = Image.open("data/example_video_dataset/ltx2/first_frame.png").convert("RGB").resize((width, height)) +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=42, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + tile_size_in_frames=80, + input_images=[first_frame], + input_images_indexes=[0], + input_images_strength=1.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_onestage_i2av_first.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py new file mode 100644 index 000000000..5c21db200 --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-I2AV-TwoStage.py @@ -0,0 +1,84 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from PIL import Image +from modelscope import dataset_snapshot_download + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], + stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), + stage2_lora_strength=1.0, +) +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +prompt = "Two cute orange cats, wearing boxing gloves, stand in a boxing ring and fight each other. They are punching each other fast and yelling: 'I will win!'" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +first_frame = Image.open("data/example_video_dataset/ltx2/first_frame.png").convert("RGB").resize((width, height)) +last_frame = Image.open("data/example_video_dataset/ltx2/last_frame.png").convert("RGB").resize((width, height)) +# first frame +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=42, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + tile_size_in_frames=80, + use_two_stage_pipeline=True, + input_images=[first_frame], + input_images_indexes=[0], + input_images_strength=1.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_twostage_i2av_first.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) +pipe.clear_lora() + +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=42, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + tile_size_in_frames=80, + use_two_stage_pipeline=True, + input_images=[first_frame, last_frame], + input_images_indexes=[0, num_frames - 1], + input_images_strength=1.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_twostage_i2av_keyframes.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py b/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py new file mode 100644 index 000000000..456ab29e6 --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py @@ -0,0 +1,69 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data import VideoData +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from modelscope import dataset_snapshot_download + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], +) +pipe.load_lora( + pipe.dit, + ModelConfig( + model_id="Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler", + origin_file_pattern="ltx-2.5-22b-ic-lora-pixel-spatial-upscaler-x2-1.0.safetensors", + ), +) + +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +reference_video = VideoData("data/example_video_dataset/ltx2/video2.mp4", height=height // 4, width=width // 4).raw_data() +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + frame_rate=24, + in_context_videos=[reference_video], + in_context_downsample_factor=2, + tiled=True, + tile_size_in_frames=80, + cfg_scale=1.0, + num_inference_steps=8, + use_distilled_pipeline=True, + use_two_stage_pipeline=True, + clear_lora_before_state_two=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_pixel_spatial_upscale.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference/LTX-2.5-T2A.py b/examples/ltx2/model_inference/LTX-2.5-T2A.py new file mode 100644 index 000000000..5a418512f --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-T2A.py @@ -0,0 +1,39 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.audio import save_audio + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), + ], +) + +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +_, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + num_frames=121, + frame_rate=24, + num_inference_steps=30, + audio_only=True, +) +save_audio(audio, pipe.audio_vocoder.output_sampling_rate, "ltx2.5_t2a.wav") diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py new file mode 100644 index 000000000..786d0ace5 --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-DistilledPipeline.py @@ -0,0 +1,58 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), + ], +) + +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width = 512 * 2, 768 * 2 +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=42, + height=height, + width=width, + frame_rate=24, + auto_duration=True, + auto_duration_min_seconds=1.0, + auto_duration_max_seconds=20.0, + cfg_scale=1.0, + num_inference_steps=8, + use_distilled_pipeline=True, + use_two_stage_pipeline=True, + tiled=True, + tile_size_in_frames=80, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_distilled_t2av.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py new file mode 100644 index 000000000..cc3f68b09 --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-INT8-ConvRot.py @@ -0,0 +1,55 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-comfy-int8-convrot.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], +) + +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + frame_rate=24, + cfg_scale=1.0, + num_inference_steps=8, + use_distilled_pipeline=True, + use_two_stage_pipeline=True, + tiled=True, + tile_size_in_frames=80, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_int8_convrot_t2av.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py new file mode 100644 index 000000000..7fba1b295 --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-OneStage.py @@ -0,0 +1,48 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ], +) +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512, 768, 121 +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + tile_size_in_frames=80, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_onestage_t2av.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py new file mode 100644 index 000000000..099a96259 --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage-Retake.py @@ -0,0 +1,69 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from diffsynth.utils.data.audio import read_audio +from diffsynth.utils.data import VideoData +from modelscope import dataset_snapshot_download + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], + stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), + stage2_lora_strength=1.0, +) + +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] + +height, width, num_frames, frame_rate = 512 * 2, 768 * 2, 121, 24 +path = "data/example_video_dataset/ltx2/video2.mp4" +video = VideoData(path, height=height, width=width).raw_data()[:num_frames] +assert len(video) == num_frames, f"Input video has {len(video)} frames, but expected {num_frames} frames based on the specified num_frames argument." +audio, audio_sample_rate = read_audio(path) + +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + retake_video=video, + retake_video_regions=[(1, 2), (3, 4)], + retake_audio=audio, + audio_sample_rate=audio_sample_rate, + retake_audio_regions=[(0, 1), (4, 5)], + seed=43, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + tiled=True, + tile_size_in_frames=80, + use_two_stage_pipeline=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_twostage_retake.mp4", + fps=frame_rate, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py new file mode 100644 index 000000000..e2234a9c7 --- /dev/null +++ b/examples/ltx2/model_inference/LTX-2.5-T2AV-TwoStage.py @@ -0,0 +1,52 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], + stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), + stage2_lora_strength=1.0, +) +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=42, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + tile_size_in_frames=80, + use_two_stage_pipeline=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_twostage_t2av.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.3-A2V-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.3-A2V-TwoStage.py index 022160bb9..1a5437174 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.3-A2V-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.3-A2V-TwoStage.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-DistilledPipeline.py b/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-DistilledPipeline.py index 5585efc67..54ee01e7c 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-DistilledPipeline.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-TwoStage.py index d23bd619d..a7e80b2cc 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.3-I2AV-TwoStage.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-DistilledPipeline.py index 8d67de931..fa3d8f2c2 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-DistilledPipeline.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-DistilledPipeline.py @@ -18,7 +18,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py index 094289a7a..e14c87a67 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Motion-Track-Control.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Union-Control.py b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Union-Control.py index b015f06ca..2f0ef7420 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Union-Control.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-IC-LoRA-Union-Control.py @@ -20,7 +20,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage-Retake.py b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage-Retake.py index 65a6ebfe5..40185f9d8 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage-Retake.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage-Retake.py @@ -21,7 +21,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage.py index a954b40e8..a37ae6b01 100644 --- a/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage.py +++ b/examples/ltx2/model_inference_low_vram/LTX-2.3-T2AV-TwoStage.py @@ -18,7 +18,7 @@ model_configs=[ ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized", origin_file_pattern="model-*.safetensors", **vram_config), ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-dev.safetensors", **vram_config), - ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", **vram_config), ], tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"), stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"), diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py new file mode 100644 index 000000000..95d3d24e9 --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-A2V-TwoStage.py @@ -0,0 +1,62 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from diffsynth.utils.data.audio import read_audio +from modelscope import dataset_snapshot_download + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], + stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), + stage2_lora_strength=1.0, + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) + +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree." +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames, frame_rate = 512 * 2, 768 * 2, 121, 24 +duration = num_frames / frame_rate +audio, audio_sample_rate = read_audio("data/example_video_dataset/ltx2/sing.MP3", start_time=1, duration=duration) +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + retake_audio=audio, + audio_sample_rate=audio_sample_rate, + seed=42, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + tiled=True, + tile_size_in_frames=80, + use_two_stage_pipeline=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_twostage_a2v.mp4", + fps=frame_rate, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py new file mode 100644 index 000000000..4c937759c --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-OneStage.py @@ -0,0 +1,56 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from PIL import Image +from modelscope import dataset_snapshot_download + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ], + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +prompt = "Two cute orange cats, wearing boxing gloves, stand in a boxing ring and fight each other. They are punching each other fast and yelling: 'I will win!'" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512, 768, 121 +first_frame = Image.open("data/example_video_dataset/ltx2/first_frame.png").convert("RGB").resize((width, height)) +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=42, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + tile_size_in_frames=80, + input_images=[first_frame], + input_images_indexes=[0], + input_images_strength=1.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_onestage_i2av_first.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py new file mode 100644 index 000000000..5d7dbe15c --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-I2AV-TwoStage.py @@ -0,0 +1,85 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from PIL import Image +from modelscope import dataset_snapshot_download + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], + stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), + stage2_lora_strength=1.0, + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +prompt = "Two cute orange cats, wearing boxing gloves, stand in a boxing ring and fight each other. They are punching each other fast and yelling: 'I will win!'" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +first_frame = Image.open("data/example_video_dataset/ltx2/first_frame.png").convert("RGB").resize((width, height)) +last_frame = Image.open("data/example_video_dataset/ltx2/last_frame.png").convert("RGB").resize((width, height)) +# first frame +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=42, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + tile_size_in_frames=80, + use_two_stage_pipeline=True, + input_images=[first_frame], + input_images_indexes=[0], + input_images_strength=1.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_twostage_i2av_first.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) +pipe.clear_lora() + +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=42, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + tile_size_in_frames=80, + use_two_stage_pipeline=True, + input_images=[first_frame, last_frame], + input_images_indexes=[0, num_frames - 1], + input_images_strength=1.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_twostage_i2av_keyframes.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py new file mode 100644 index 000000000..56a7ce198 --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-IC-LoRA-Pixel-Spatial-Upscaler.py @@ -0,0 +1,70 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data import VideoData +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from modelscope import dataset_snapshot_download + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +pipe.load_lora( + pipe.dit, + ModelConfig( + model_id="Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler", + origin_file_pattern="ltx-2.5-22b-ic-lora-pixel-spatial-upscaler-x2-1.0.safetensors", + ), +) + +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +reference_video = VideoData("data/example_video_dataset/ltx2/video2.mp4", height=height // 4, width=width // 4).raw_data() +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + frame_rate=24, + in_context_videos=[reference_video], + in_context_downsample_factor=2, + tiled=True, + tile_size_in_frames=80, + cfg_scale=1.0, + num_inference_steps=8, + use_distilled_pipeline=True, + use_two_stage_pipeline=True, + clear_lora_before_state_two=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_pixel_spatial_upscale.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py new file mode 100644 index 000000000..c7fa4c517 --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2A.py @@ -0,0 +1,40 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.audio import save_audio + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), + ], + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) + +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +_, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + num_frames=121, + frame_rate=24, + num_inference_steps=30, + audio_only=True, +) +save_audio(audio, pipe.audio_vocoder.output_sampling_rate, "ltx2.5_t2a.wav") diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py new file mode 100644 index 000000000..80af71a7f --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-DistilledPipeline.py @@ -0,0 +1,59 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="model_patches/ltx-2.5-duration-head-bf16.safetensors", **vram_config), + ], + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) + +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width = 512 * 2, 768 * 2 +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=42, + height=height, + width=width, + frame_rate=24, + auto_duration=True, + auto_duration_min_seconds=1.0, + auto_duration_max_seconds=20.0, + cfg_scale=1.0, + num_inference_steps=8, + use_distilled_pipeline=True, + use_two_stage_pipeline=True, + tiled=True, + tile_size_in_frames=80, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_distilled_t2av.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py new file mode 100644 index 000000000..d144681a4 --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-INT8-ConvRot.py @@ -0,0 +1,56 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-distilled-transformer-comfy-int8-convrot.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) + +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + frame_rate=24, + cfg_scale=1.0, + num_inference_steps=8, + use_distilled_pipeline=True, + use_two_stage_pipeline=True, + tiled=True, + tile_size_in_frames=80, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_int8_convrot_t2av.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py new file mode 100644 index 000000000..7e5227d02 --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-OneStage.py @@ -0,0 +1,49 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ], + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512, 768, 121 +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + tile_size_in_frames=80, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_onestage_t2av.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py new file mode 100644 index 000000000..fffae6d2b --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage-Retake.py @@ -0,0 +1,70 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from diffsynth.utils.data.audio import read_audio +from diffsynth.utils.data import VideoData +from modelscope import dataset_snapshot_download + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], + stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), + stage2_lora_strength=1.0, + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) + +dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset") +prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree. She sings: 'Mummy don't know daddy's getting hot. At the body shop'" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] + +height, width, num_frames, frame_rate = 512 * 2, 768 * 2, 121, 24 +path = "data/example_video_dataset/ltx2/video2.mp4" +video = VideoData(path, height=height, width=width).raw_data()[:num_frames] +assert len(video) == num_frames, f"Input video has {len(video)} frames, but expected {num_frames} frames based on the specified num_frames argument." +audio, audio_sample_rate = read_audio(path) + +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + retake_video=video, + retake_video_regions=[(1, 2), (3, 4)], + retake_audio=audio, + audio_sample_rate=audio_sample_rate, + retake_audio_regions=[(0, 1), (4, 5)], + seed=43, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + tiled=True, + tile_size_in_frames=80, + use_two_stage_pipeline=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_twostage_retake.mp4", + fps=frame_rate, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py new file mode 100644 index 000000000..be12094d4 --- /dev/null +++ b/examples/ltx2/model_inference_low_vram/LTX-2.5-T2AV-TwoStage.py @@ -0,0 +1,53 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": "disk", + "onload_device": "disk", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + # For lower VRAM and faster decoding, replace the line above with the conv vae decoder: + # ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-conv-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", **vram_config), + ], + stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors"), + stage2_lora_strength=1.0, + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”" +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512 * 2, 768 * 2, 121 +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=42, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + tile_size_in_frames=80, + use_two_stage_pipeline=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_twostage_t2av.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_training/full/LTX-2.5-I2AV-splited.sh b/examples/ltx2/model_training/full/LTX-2.5-I2AV-splited.sh new file mode 100644 index 000000000..00ccf3fc8 --- /dev/null +++ b/examples/ltx2/model_training/full/LTX-2.5-I2AV-splited.sh @@ -0,0 +1,41 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ltx2/LTX-2.5-I2AV-splited/*" --local_dir ./data/diffsynth_example_dataset + +# Splited Training +accelerate launch examples/ltx2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/ltx2/LTX-2.5-I2AV-splited \ + --dataset_metadata_path data/diffsynth_example_dataset/ltx2/LTX-2.5-I2AV-splited/metadata.csv \ + --data_file_keys "video,input_audio" \ + --extra_inputs "input_audio,input_image" \ + --height 512 \ + --width 768 \ + --num_frames 121 \ + --dataset_repeat 1 \ + --tokenizer_path "./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/LTX2.5-I2AV-full-splited-cache" \ + --trainable_models "dit" \ + --use_gradient_checkpointing \ + --task "sft:data_process" + +accelerate launch --config_file examples/ltx2/model_training/full/accelerate_config_zero2offload.yaml examples/ltx2/model_training/train.py \ + --dataset_base_path ./models/train/LTX2.5-I2AV-full-splited-cache \ + --data_file_keys "video,input_audio" \ + --extra_inputs "input_audio,input_image" \ + --height 512 \ + --width 768 \ + --num_frames 121 \ + --dataset_repeat 100 \ + --tokenizer_path "./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ + --fp8_models "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/LTX2.5-I2AV-full" \ + --trainable_models "dit" \ + --use_gradient_checkpointing \ + --find_unused_parameters \ + --task "sft:train" diff --git a/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh b/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh new file mode 100644 index 000000000..6cb69b7c5 --- /dev/null +++ b/examples/ltx2/model_training/full/LTX-2.5-T2AV-splited.sh @@ -0,0 +1,41 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ltx2/LTX-2.5-T2AV-splited/*" --local_dir ./data/diffsynth_example_dataset + +# Splited Training +accelerate launch examples/ltx2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/ltx2/LTX-2.5-T2AV-splited \ + --dataset_metadata_path data/diffsynth_example_dataset/ltx2/LTX-2.5-T2AV-splited/metadata.csv \ + --data_file_keys "video,input_audio" \ + --extra_inputs "input_audio" \ + --height 512 \ + --width 768 \ + --num_frames 121 \ + --dataset_repeat 1 \ + --tokenizer_path "./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/LTX2.5-T2AV-full-splited-cache" \ + --trainable_models "dit" \ + --use_gradient_checkpointing \ + --task "sft:data_process" + +accelerate launch --config_file examples/ltx2/model_training/full/accelerate_config_zero2offload.yaml examples/ltx2/model_training/train.py \ + --dataset_base_path ./models/train/LTX2.5-T2AV-full-splited-cache \ + --data_file_keys "video,input_audio" \ + --extra_inputs "input_audio" \ + --height 512 \ + --width 768 \ + --num_frames 121 \ + --dataset_repeat 100 \ + --tokenizer_path "./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ + --fp8_models "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 2 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/LTX2.5-T2AV-full" \ + --trainable_models "dit" \ + --use_gradient_checkpointing \ + --find_unused_parameters \ + --task "sft:train" diff --git a/examples/ltx2/model_training/full/accelerate_config_zero2offload.yaml b/examples/ltx2/model_training/full/accelerate_config_zero2offload.yaml new file mode 100644 index 000000000..8a75f3d91 --- /dev/null +++ b/examples/ltx2/model_training/full/accelerate_config_zero2offload.yaml @@ -0,0 +1,22 @@ +compute_environment: LOCAL_MACHINE +debug: false +deepspeed_config: + gradient_accumulation_steps: 1 + offload_optimizer_device: 'cpu' + offload_param_device: 'cpu' + zero3_init_flag: false + zero_stage: 2 +distributed_type: DEEPSPEED +downcast_bf16: 'no' +enable_cpu_affinity: false +machine_rank: 0 +main_training_function: main +mixed_precision: bf16 +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false diff --git a/examples/ltx2/model_training/lora/LTX-2.5-I2AV-splited.sh b/examples/ltx2/model_training/lora/LTX-2.5-I2AV-splited.sh new file mode 100644 index 000000000..3f53b3a02 --- /dev/null +++ b/examples/ltx2/model_training/lora/LTX-2.5-I2AV-splited.sh @@ -0,0 +1,44 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ltx2/LTX-2.5-I2AV-splited/*" --local_dir ./data/diffsynth_example_dataset + +# Splited Training +accelerate launch examples/ltx2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/ltx2/LTX-2.5-I2AV-splited \ + --dataset_metadata_path data/diffsynth_example_dataset/ltx2/LTX-2.5-I2AV-splited/metadata.csv \ + --data_file_keys "video,input_audio" \ + --extra_inputs "input_audio,input_image" \ + --height 512 \ + --width 768 \ + --num_frames 121 \ + --dataset_repeat 1 \ + --tokenizer_path "./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/LTX2.5-I2AV_lora-splited-cache" \ + --lora_base_model "dit" \ + --lora_target_modules "to_k,to_q,to_v,to_out.0" \ + --lora_rank 32 \ + --use_gradient_checkpointing \ + --task "sft:data_process" + +accelerate launch examples/ltx2/model_training/train.py \ + --dataset_base_path ./models/train/LTX2.5-I2AV_lora-splited-cache \ + --data_file_keys "video,input_audio" \ + --extra_inputs "input_audio,input_image" \ + --height 512 \ + --width 768 \ + --num_frames 121 \ + --dataset_repeat 100 \ + --tokenizer_path "./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ + --fp8_models "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/LTX2.5-I2AV_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "to_k,to_q,to_v,to_out.0" \ + --lora_rank 32 \ + --use_gradient_checkpointing \ + --task "sft:train" diff --git a/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh b/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh new file mode 100644 index 000000000..af873ff64 --- /dev/null +++ b/examples/ltx2/model_training/lora/LTX-2.5-T2AV-splited.sh @@ -0,0 +1,44 @@ +modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "ltx2/LTX-2.5-T2AV-splited/*" --local_dir ./data/diffsynth_example_dataset + +# Splited Training +accelerate launch examples/ltx2/model_training/train.py \ + --dataset_base_path data/diffsynth_example_dataset/ltx2/LTX-2.5-T2AV-splited \ + --dataset_metadata_path data/diffsynth_example_dataset/ltx2/LTX-2.5-T2AV-splited/metadata.csv \ + --data_file_keys "video,input_audio" \ + --extra_inputs "input_audio" \ + --height 512 \ + --width 768 \ + --num_frames 121 \ + --dataset_repeat 1 \ + --tokenizer_path "./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/LTX2.5-T2AV_lora-splited-cache" \ + --lora_base_model "dit" \ + --lora_target_modules "to_k,to_q,to_v,to_out.0" \ + --lora_rank 32 \ + --use_gradient_checkpointing \ + --task "sft:data_process" + +accelerate launch examples/ltx2/model_training/train.py \ + --dataset_base_path ./models/train/LTX2.5-T2AV_lora-splited-cache \ + --data_file_keys "video,input_audio" \ + --extra_inputs "input_audio" \ + --height 512 \ + --width 768 \ + --num_frames 121 \ + --dataset_repeat 100 \ + --tokenizer_path "./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" \ + --model_id_with_origin_paths "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ + --fp8_models "Lightricks/LTX-2.5:text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-video-vae-bf16.safetensors,Lightricks/LTX-2.5:vae/ltx-2.5-audio-vae-bf16.safetensors,DiffSynth-Studio/LTX-2.5-Repackage:text_encoder_post_modules.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/LTX2.5-T2AV_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "to_k,to_q,to_v,to_out.0" \ + --lora_rank 32 \ + --use_gradient_checkpointing \ + --task "sft:train" diff --git a/examples/ltx2/model_training/scripts/split_model_statedicts_ltx2.5.py b/examples/ltx2/model_training/scripts/split_model_statedicts_ltx2.5.py new file mode 100644 index 000000000..c0d0bf22d --- /dev/null +++ b/examples/ltx2/model_training/scripts/split_model_statedicts_ltx2.5.py @@ -0,0 +1,51 @@ +import os + +import numpy as np + +from safetensors import safe_open +from safetensors.torch import save_file + +from diffsynth import hash_state_dict_keys +from diffsynth.core import load_state_dict +from diffsynth.models.model_loader import ModelPool + +model_pool = ModelPool() +os.makedirs("models/DiffSynth-Studio/LTX-2.5-Repackage", exist_ok=True) + +SOURCES = ( + "models/Lightricks/LTX-2.5/text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", + "models/Lightricks/LTX-2.5/diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", +) + +def target_name(name): + if name.startswith("text_embedding_projection."): + return "feature_extractor." + name.removeprefix("text_embedding_projection.") + if name.startswith("model.diffusion_model.video_embeddings_connector."): + return "connectors.video_connector." + name.removeprefix("model.diffusion_model.video_embeddings_connector.") + if name.startswith("model.diffusion_model.audio_embeddings_connector."): + return "connectors.audio_connector." + name.removeprefix("model.diffusion_model.audio_embeddings_connector.") + return None + + +text_encoder_post_modules_state_dict = {} +for path in SOURCES: + with safe_open(path, framework="pt", device="cpu") as handle: + for name in handle.keys(): + new_name = target_name(name) + if new_name is not None: + text_encoder_post_modules_state_dict[new_name] = handle.get_tensor(name) + +save_file(text_encoder_post_modules_state_dict, "models/DiffSynth-Studio/LTX-2.5-Repackage/text_encoder_post_modules.safetensors") +print(f"text_encoder_post_modules keys hash: {hash_state_dict_keys(text_encoder_post_modules_state_dict)}") +model_pool.auto_load_model("models/DiffSynth-Studio/LTX-2.5-Repackage/text_encoder_post_modules.safetensors") + +tokenizer_dir = "models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer" +os.makedirs(tokenizer_dir, exist_ok=True) +with safe_open(SOURCES[0], framework="pt", device="cpu") as handle: + tokenizer_bytes = handle.get_tensor("tokenizer_json").detach().cpu().numpy().astype(np.uint8).tobytes() + config_bytes = handle.get_tensor("hf_asset__tokenizer_config.json").detach().cpu().numpy().astype(np.uint8).tobytes() +with open(os.path.join(tokenizer_dir, "tokenizer.json"), "wb") as f: + f.write(tokenizer_bytes) +with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "wb") as f: + f.write(config_bytes) +print(f"tokenizer assets written to {tokenizer_dir}") diff --git a/examples/ltx2/model_training/validate_full/LTX-2.5-I2AV.py b/examples/ltx2/model_training/validate_full/LTX-2.5-I2AV.py new file mode 100644 index 000000000..0f5afe452 --- /dev/null +++ b/examples/ltx2/model_training/validate_full/LTX-2.5-I2AV.py @@ -0,0 +1,51 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from diffsynth.utils.data import VideoData + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(path="./models/train/LTX2.5-I2AV-full/epoch-1.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ], +) +prompt = "A beautiful sunset over the ocean." +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512, 768, 121 +image = VideoData("data/diffsynth_example_dataset/ltx2/LTX-2.5-I2AV-splited/video.mp4", height=height, width=width)[0] +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=42, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + tile_size_in_frames=80, + input_images=[image], + input_images_indexes=[0], + input_images_strength=1.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_i2av_full.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py b/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py new file mode 100644 index 000000000..0d338e9f0 --- /dev/null +++ b/examples/ltx2/model_training/validate_full/LTX-2.5-T2AV.py @@ -0,0 +1,45 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(path="./models/train/LTX2.5-T2AV-full/epoch-1.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ], +) +prompt = "A beautiful sunset over the ocean." +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512, 768, 121 +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + tiled=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_t2av_full.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_training/validate_lora/LTX-2.5-I2AV.py b/examples/ltx2/model_training/validate_lora/LTX-2.5-I2AV.py new file mode 100644 index 000000000..50296faf6 --- /dev/null +++ b/examples/ltx2/model_training/validate_lora/LTX-2.5-I2AV.py @@ -0,0 +1,52 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 +from diffsynth.utils.data import VideoData + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ], +) +pipe.load_lora(pipe.dit, "models/train/LTX2.5-I2AV_lora/epoch-4.safetensors") +prompt = "A beautiful sunset over the ocean." +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512, 768, 121 +image = VideoData("data/diffsynth_example_dataset/ltx2/LTX-2.5-I2AV-splited/video.mp4", height=height, width=width)[0] +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=42, + height=height, + width=width, + num_frames=num_frames, + tiled=True, + tile_size_in_frames=80, + input_images=[image], + input_images_indexes=[0], + input_images_strength=1.0, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_i2av_lora.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +) diff --git a/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py b/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py new file mode 100644 index 000000000..d2932fcba --- /dev/null +++ b/examples/ltx2/model_training/validate_lora/LTX-2.5-T2AV.py @@ -0,0 +1,46 @@ +import torch +from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig +from diffsynth.utils.data.media_io_ltx2 import write_video_audio_ltx2 + +vram_config = { + "offload_dtype": torch.bfloat16, + "offload_device": "cpu", + "onload_dtype": torch.bfloat16, + "onload_device": "cuda", + "preparing_dtype": torch.bfloat16, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = LTX2AudioVideoPipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + tokenizer_config=ModelConfig(path="./models/DiffSynth-Studio/LTX-2.5-Repackage/tokenizer"), + model_configs=[ + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", **vram_config), + ModelConfig(model_id="DiffSynth-Studio/LTX-2.5-Repackage", origin_file_pattern="text_encoder_post_modules.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-video-vae-bf16.safetensors", **vram_config), + ModelConfig(model_id="Lightricks/LTX-2.5", origin_file_pattern="vae/ltx-2.5-audio-vae-bf16.safetensors", **vram_config), + ], +) +pipe.load_lora(pipe.dit, "models/train/LTX2.5-T2AV_lora/epoch-4.safetensors") +prompt = "A beautiful sunset over the ocean." +negative_prompt = pipe.default_negative_prompt["LTX-2.5"] +height, width, num_frames = 512, 768, 121 +video, audio = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + seed=43, + height=height, + width=width, + num_frames=num_frames, + tiled=True, +) +write_video_audio_ltx2( + video=video, + audio=audio, + output_path="ltx2.5_t2av_lora.mp4", + fps=24, + audio_sample_rate=pipe.audio_vocoder.output_sampling_rate, +)