github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add support for training LTX 2.3 (#745)

Browse Source 
* Initial support for ltx 2.3. Still needs a lot of testing to make sure it is all right.

* bump version

* Handle lora renaming keys for new ltx 2.3 layers
This commit is contained in:
Jaret Burkett
2026-03-22 17:56:59 -06:00
committed by GitHub
parent 7e1fdc3844
commit 859635e95b
8 changed files with 707 additions and 236 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
dgx_requirements.txt
View File

@@ -13,7 +13,7 @@ opencv_python_headless==4.11.0.86
# additional dependencies matching requirements.txt:
torchao==0.10.0
safetensors
git+https://github.com/huggingface/diffusers@8600b4c10d67b0ce200f664204358747bd53c775
git+https://github.com/huggingface/diffusers@072d15ee4289ffdc3aa9d65f8b94bc9271319d21
transformers==4.57.3
lycoris-lora==1.8.3
flatten_json

3
extensions_built_in/diffusion_models/__init__.py
View File

@@ -7,7 +7,7 @@ from .wan22 import Wan225bModel, Wan2214bModel, Wan2214bI2VModel
from .qwen_image import QwenImageModel, QwenImageEditModel, QwenImageEditPlusModel
from .flux2 import Flux2Model, Flux2Klein4BModel, Flux2Klein9BModel
from .z_image import ZImageModel
from .ltx2 import LTX2Model
from .ltx2 import LTX2Model, LTX23Model
from .zeta_chroma import ZetaChromaModel
AI_TOOLKIT_MODELS = [
@@ -28,6 +28,7 @@ AI_TOOLKIT_MODELS = [
Flux2Model,
ZImageModel,
LTX2Model,
LTX23Model,
Flux2Klein4BModel,
Flux2Klein9BModel,
ZetaChromaModel,

2
extensions_built_in/diffusion_models/ltx2/__init__.py
View File

@@ -1 +1 @@
from .ltx2 import LTX2Model
from .ltx2 import LTX2Model, LTX23Model

778
extensions_built_in/diffusion_models/ltx2/convert_ltx2_to_diffusers.py
View File

@@ -1,4 +1,4 @@
# ref https://github.com/huggingface/diffusers/blob/249ae1f853be8775c7d0b3a26ca4fbc4f6aa9998/scripts/convert_ltx2_to_diffusers.py
# ref https://github.com/huggingface/diffusers/blob/17b53f08661732caca6a546295950fc4b1696ad7/scripts/convert_ltx2_to_diffusers.py
from contextlib import nullcontext
from typing import Any, Dict, Tuple
@@ -11,7 +11,7 @@ from diffusers import (
AutoencoderKLLTX2Video,
LTX2VideoTransformer3DModel,
)
from diffusers.pipelines.ltx2 import LTX2TextConnectors, LTX2Vocoder
from diffusers.pipelines.ltx2 import LTX2TextConnectors, LTX2Vocoder, LTX2VocoderWithBWE
from diffusers.utils.import_utils import is_accelerate_available
@@ -38,6 +38,12 @@ LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT = {
"k_norm": "norm_k",
}
LTX_2_3_TRANSFORMER_KEYS_RENAME_DICT = {
**LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT,
"audio_prompt_adaln_single": "audio_prompt_adaln",
"prompt_adaln_single": "prompt_adaln",
}
LTX_2_0_VIDEO_VAE_RENAME_DICT = {
# Encoder
"down_blocks.0": "down_blocks.0",
@@ -57,6 +63,8 @@ LTX_2_0_VIDEO_VAE_RENAME_DICT = {
"up_blocks.4": "up_blocks.1",
"up_blocks.5": "up_blocks.2.upsamplers.0",
"up_blocks.6": "up_blocks.2",
"last_time_embedder": "time_embedder",
"last_scale_shift_table": "scale_shift_table",
# Common
# For all 3D ResNets
"res_blocks": "resnets",
@@ -64,6 +72,13 @@ LTX_2_0_VIDEO_VAE_RENAME_DICT = {
"per_channel_statistics.std-of-means": "latents_std",
}
LTX_2_3_VIDEO_VAE_RENAME_DICT = {
**LTX_2_0_VIDEO_VAE_RENAME_DICT,
# Decoder extra blocks
"up_blocks.7": "up_blocks.3.upsamplers.0",
"up_blocks.8": "up_blocks.3",
}
LTX_2_0_AUDIO_VAE_RENAME_DICT = {
"per_channel_statistics.mean-of-means": "latents_mean",
"per_channel_statistics.std-of-means": "latents_std",
@@ -76,10 +91,34 @@ LTX_2_0_VOCODER_RENAME_DICT = {
"conv_post": "conv_out",
}
LTX_2_0_TEXT_ENCODER_RENAME_DICT = {
LTX_2_3_VOCODER_RENAME_DICT = {
# Handle upsamplers ("ups" --> "upsamplers") due to name clash
"resblocks": "resnets",
"conv_pre": "conv_in",
"conv_post": "conv_out",
"act_post": "act_out",
"downsample.lowpass": "downsample",
}
LTX_2_0_CONNECTORS_KEYS_RENAME_DICT = {
"connectors.": "",
"video_embeddings_connector": "video_connector",
"audio_embeddings_connector": "audio_connector",
"transformer_1d_blocks": "transformer_blocks",
"text_embedding_projection.aggregate_embed": "text_proj_in",
# Attention QK Norms
"q_norm": "norm_q",
"k_norm": "norm_k",
}
LTX_2_3_CONNECTORS_KEYS_RENAME_DICT = {
"connectors.": "",
"video_embeddings_connector": "video_connector",
"audio_embeddings_connector": "audio_connector",
"transformer_1d_blocks": "transformer_blocks",
# LTX-2.3 uses per-modality embedding projections
"text_embedding_projection.audio_aggregate_embed": "audio_text_proj_in",
"text_embedding_projection.video_aggregate_embed": "video_text_proj_in",
# Attention QK Norms
"q_norm": "norm_q",
"k_norm": "norm_k",
@@ -125,23 +164,24 @@ def convert_ltx2_audio_vae_per_channel_statistics(
return
def convert_ltx2_3_vocoder_upsamplers(key: str, state_dict: dict[str, Any]) -> None:
# Skip if not a weight, bias
if ".weight" not in key and ".bias" not in key:
return
if ".ups." in key:
new_key = key.replace(".ups.", ".upsamplers.")
param = state_dict.pop(key)
state_dict[new_key] = param
return
LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP = {
"video_embeddings_connector": remove_keys_inplace,
"audio_embeddings_connector": remove_keys_inplace,
"adaln_single": convert_ltx2_transformer_adaln_single,
}
LTX_2_0_CONNECTORS_KEYS_RENAME_DICT = {
"connectors.": "",
"video_embeddings_connector": "video_connector",
"audio_embeddings_connector": "audio_connector",
"transformer_1d_blocks": "transformer_blocks",
"text_embedding_projection.aggregate_embed": "text_proj_in",
# Attention QK Norms
"q_norm": "norm_q",
"k_norm": "norm_k",
}
LTX_2_0_VAE_SPECIAL_KEYS_REMAP = {
"per_channel_statistics.channel": remove_keys_inplace,
"per_channel_statistics.mean-of-stds": remove_keys_inplace,
@@ -151,6 +191,12 @@ LTX_2_0_AUDIO_VAE_SPECIAL_KEYS_REMAP = {}
LTX_2_0_VOCODER_SPECIAL_KEYS_REMAP = {}
LTX_2_3_VOCODER_SPECIAL_KEYS_REMAP = {
".ups.": convert_ltx2_3_vocoder_upsamplers,
}
LTX_2_0_CONNECTORS_SPECIAL_KEYS_REMAP = {}
def split_transformer_and_connector_state_dict(
state_dict: Dict[str, Any],
@@ -159,7 +205,7 @@ def split_transformer_and_connector_state_dict(
"video_embeddings_connector",
"audio_embeddings_connector",
"transformer_1d_blocks",
"text_embedding_projection.aggregate_embed",
"text_embedding_projection",
"connectors.",
"video_connector",
"audio_connector",
@@ -176,87 +222,243 @@ def split_transformer_and_connector_state_dict(
return transformer_state_dict, connector_state_dict
def get_ltx2_transformer_config() -> Tuple[
Dict[str, Any], Dict[str, Any], Dict[str, Any]
]:
config = {
"model_id": "diffusers-internal-dev/new-ltx-model",
"diffusers_config": {
"in_channels": 128,
"out_channels": 128,
"patch_size": 1,
"patch_size_t": 1,
"num_attention_heads": 32,
"attention_head_dim": 128,
"cross_attention_dim": 4096,
"vae_scale_factors": (8, 32, 32),
"pos_embed_max_pos": 20,
"base_height": 2048,
"base_width": 2048,
"audio_in_channels": 128,
"audio_out_channels": 128,
"audio_patch_size": 1,
"audio_patch_size_t": 1,
"audio_num_attention_heads": 32,
"audio_attention_head_dim": 64,
"audio_cross_attention_dim": 2048,
"audio_scale_factor": 4,
"audio_pos_embed_max_pos": 20,
"audio_sampling_rate": 16000,
"audio_hop_length": 160,
"num_layers": 48,
"activation_fn": "gelu-approximate",
"qk_norm": "rms_norm_across_heads",
"norm_elementwise_affine": False,
"norm_eps": 1e-6,
"caption_channels": 3840,
"attention_bias": True,
"attention_out_bias": True,
"rope_theta": 10000.0,
"rope_double_precision": True,
"causal_offset": 1,
"timestep_scale_multiplier": 1000,
"cross_attn_timestep_scale_multiplier": 1000,
"rope_type": "split",
},
}
rename_dict = LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT
special_keys_remap = LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP
def get_ltx2_transformer_config(
version: str = "2.0",
) -> Tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]:
if version == "test":
# Produces a transformer of the same size as used in test_models_transformer_ltx2.py
config = {
"model_id": "diffusers-internal-dev/dummy-ltx2",
"diffusers_config": {
"in_channels": 4,
"out_channels": 4,
"patch_size": 1,
"patch_size_t": 1,
"num_attention_heads": 2,
"attention_head_dim": 8,
"cross_attention_dim": 16,
"vae_scale_factors": (8, 32, 32),
"pos_embed_max_pos": 20,
"base_height": 2048,
"base_width": 2048,
"audio_in_channels": 4,
"audio_out_channels": 4,
"audio_patch_size": 1,
"audio_patch_size_t": 1,
"audio_num_attention_heads": 2,
"audio_attention_head_dim": 4,
"audio_cross_attention_dim": 8,
"audio_scale_factor": 4,
"audio_pos_embed_max_pos": 20,
"audio_sampling_rate": 16000,
"audio_hop_length": 160,
"num_layers": 2,
"activation_fn": "gelu-approximate",
"qk_norm": "rms_norm_across_heads",
"norm_elementwise_affine": False,
"norm_eps": 1e-6,
"caption_channels": 16,
"attention_bias": True,
"attention_out_bias": True,
"rope_theta": 10000.0,
"rope_double_precision": False,
"causal_offset": 1,
"timestep_scale_multiplier": 1000,
"cross_attn_timestep_scale_multiplier": 1,
},
}
rename_dict = LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT
special_keys_remap = LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP
elif version == "2.0":
config = {
"model_id": "Lightricks/LTX-2",
"diffusers_config": {
"in_channels": 128,
"out_channels": 128,
"patch_size": 1,
"patch_size_t": 1,
"num_attention_heads": 32,
"attention_head_dim": 128,
"cross_attention_dim": 4096,
"vae_scale_factors": (8, 32, 32),
"pos_embed_max_pos": 20,
"base_height": 2048,
"base_width": 2048,
"gated_attn": False,
"cross_attn_mod": False,
"audio_in_channels": 128,
"audio_out_channels": 128,
"audio_patch_size": 1,
"audio_patch_size_t": 1,
"audio_num_attention_heads": 32,
"audio_attention_head_dim": 64,
"audio_cross_attention_dim": 2048,
"audio_scale_factor": 4,
"audio_pos_embed_max_pos": 20,
"audio_sampling_rate": 16000,
"audio_hop_length": 160,
"audio_gated_attn": False,
"audio_cross_attn_mod": False,
"num_layers": 48,
"activation_fn": "gelu-approximate",
"qk_norm": "rms_norm_across_heads",
"norm_elementwise_affine": False,
"norm_eps": 1e-6,
"caption_channels": 3840,
"attention_bias": True,
"attention_out_bias": True,
"rope_theta": 10000.0,
"rope_double_precision": True,
"causal_offset": 1,
"timestep_scale_multiplier": 1000,
"cross_attn_timestep_scale_multiplier": 1000,
"rope_type": "split",
"use_prompt_embeddings": True,
"perturbed_attn": False,
},
}
rename_dict = LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT
special_keys_remap = LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP
elif version == "2.3":
config = {
"model_id": "Lightricks/LTX-2.3",
"diffusers_config": {
"in_channels": 128,
"out_channels": 128,
"patch_size": 1,
"patch_size_t": 1,
"num_attention_heads": 32,
"attention_head_dim": 128,
"cross_attention_dim": 4096,
"vae_scale_factors": (8, 32, 32),
"pos_embed_max_pos": 20,
"base_height": 2048,
"base_width": 2048,
"gated_attn": True,
"cross_attn_mod": True,
"audio_in_channels": 128,
"audio_out_channels": 128,
"audio_patch_size": 1,
"audio_patch_size_t": 1,
"audio_num_attention_heads": 32,
"audio_attention_head_dim": 64,
"audio_cross_attention_dim": 2048,
"audio_scale_factor": 4,
"audio_pos_embed_max_pos": 20,
"audio_sampling_rate": 16000,
"audio_hop_length": 160,
"audio_gated_attn": True,
"audio_cross_attn_mod": True,
"num_layers": 48,
"activation_fn": "gelu-approximate",
"qk_norm": "rms_norm_across_heads",
"norm_elementwise_affine": False,
"norm_eps": 1e-6,
"caption_channels": 3840,
"attention_bias": True,
"attention_out_bias": True,
"rope_theta": 10000.0,
"rope_double_precision": True,
"causal_offset": 1,
"timestep_scale_multiplier": 1000,
"cross_attn_timestep_scale_multiplier": 1000,
"rope_type": "split",
"use_prompt_embeddings": False,
"perturbed_attn": True,
},
}
rename_dict = LTX_2_3_TRANSFORMER_KEYS_RENAME_DICT
special_keys_remap = LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP
return config, rename_dict, special_keys_remap
def get_ltx2_connectors_config() -> Tuple[
Dict[str, Any], Dict[str, Any], Dict[str, Any]
]:
config = {
"model_id": "diffusers-internal-dev/new-ltx-model",
"diffusers_config": {
"caption_channels": 3840,
"text_proj_in_factor": 49,
"video_connector_num_attention_heads": 30,
"video_connector_attention_head_dim": 128,
"video_connector_num_layers": 2,
"video_connector_num_learnable_registers": 128,
"audio_connector_num_attention_heads": 30,
"audio_connector_attention_head_dim": 128,
"audio_connector_num_layers": 2,
"audio_connector_num_learnable_registers": 128,
"connector_rope_base_seq_len": 4096,
"rope_theta": 10000.0,
"rope_double_precision": True,
"causal_temporal_positioning": False,
"rope_type": "split",
},
}
rename_dict = LTX_2_0_CONNECTORS_KEYS_RENAME_DICT
special_keys_remap = {}
def get_ltx2_connectors_config(
version: str = "2.0",
) -> tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]:
if version == "test":
config = {
"model_id": "diffusers-internal-dev/dummy-ltx2",
"diffusers_config": {
"caption_channels": 16,
"text_proj_in_factor": 3,
"video_connector_num_attention_heads": 4,
"video_connector_attention_head_dim": 8,
"video_connector_num_layers": 1,
"video_connector_num_learnable_registers": None,
"audio_connector_num_attention_heads": 4,
"audio_connector_attention_head_dim": 8,
"audio_connector_num_layers": 1,
"audio_connector_num_learnable_registers": None,
"connector_rope_base_seq_len": 32,
"rope_theta": 10000.0,
"rope_double_precision": False,
"causal_temporal_positioning": False,
},
}
elif version == "2.0":
config = {
"model_id": "Lightricks/LTX-2",
"diffusers_config": {
"caption_channels": 3840,
"text_proj_in_factor": 49,
"video_connector_num_attention_heads": 30,
"video_connector_attention_head_dim": 128,
"video_connector_num_layers": 2,
"video_connector_num_learnable_registers": 128,
"video_gated_attn": False,
"audio_connector_num_attention_heads": 30,
"audio_connector_attention_head_dim": 128,
"audio_connector_num_layers": 2,
"audio_connector_num_learnable_registers": 128,
"audio_gated_attn": False,
"connector_rope_base_seq_len": 4096,
"rope_theta": 10000.0,
"rope_double_precision": True,
"causal_temporal_positioning": False,
"rope_type": "split",
"per_modality_projections": False,
"proj_bias": False,
},
}
rename_dict = LTX_2_0_CONNECTORS_KEYS_RENAME_DICT
special_keys_remap = LTX_2_0_CONNECTORS_SPECIAL_KEYS_REMAP
elif version == "2.3":
config = {
"model_id": "Lightricks/LTX-2.3",
"diffusers_config": {
"caption_channels": 3840,
"text_proj_in_factor": 49,
"video_connector_num_attention_heads": 32,
"video_connector_attention_head_dim": 128,
"video_connector_num_layers": 8,
"video_connector_num_learnable_registers": 128,
"video_gated_attn": True,
"audio_connector_num_attention_heads": 32,
"audio_connector_attention_head_dim": 64,
"audio_connector_num_layers": 8,
"audio_connector_num_learnable_registers": 128,
"audio_gated_attn": True,
"connector_rope_base_seq_len": 4096,
"rope_theta": 10000.0,
"rope_double_precision": True,
"causal_temporal_positioning": False,
"rope_type": "split",
"per_modality_projections": True,
"video_hidden_dim": 4096,
"audio_hidden_dim": 2048,
"proj_bias": True,
},
}
rename_dict = LTX_2_3_CONNECTORS_KEYS_RENAME_DICT
special_keys_remap = LTX_2_0_CONNECTORS_SPECIAL_KEYS_REMAP
return config, rename_dict, special_keys_remap
def convert_ltx2_transformer(original_state_dict: Dict[str, Any]) -> Dict[str, Any]:
config, rename_dict, special_keys_remap = get_ltx2_transformer_config()
def convert_ltx2_transformer(
original_state_dict: Dict[str, Any], version: str = "2.0"
) -> Dict[str, Any]:
config, rename_dict, special_keys_remap = get_ltx2_transformer_config(version)
diffusers_config = config["diffusers_config"]
transformer_state_dict, _ = split_transformer_and_connector_state_dict(
@@ -285,8 +487,10 @@ def convert_ltx2_transformer(original_state_dict: Dict[str, Any]) -> Dict[str, A
return transformer
def convert_ltx2_connectors(original_state_dict: Dict[str, Any]) -> LTX2TextConnectors:
config, rename_dict, special_keys_remap = get_ltx2_connectors_config()
def convert_ltx2_connectors(
original_state_dict: Dict[str, Any], version: str = "2.0"
) -> LTX2TextConnectors:
config, rename_dict, special_keys_remap = get_ltx2_connectors_config(version)
diffusers_config = config["diffusers_config"]
_, connector_state_dict = split_transformer_and_connector_state_dict(
@@ -314,55 +518,153 @@ def convert_ltx2_connectors(original_state_dict: Dict[str, Any]) -> LTX2TextConn
return connectors
def get_ltx2_video_vae_config() -> Tuple[
Dict[str, Any], Dict[str, Any], Dict[str, Any]
]:
config = {
"model_id": "diffusers-internal-dev/dummy-ltx2",
"diffusers_config": {
"in_channels": 3,
"out_channels": 3,
"latent_channels": 128,
"block_out_channels": (256, 512, 1024, 2048),
"down_block_types": (
"LTX2VideoDownBlock3D",
"LTX2VideoDownBlock3D",
"LTX2VideoDownBlock3D",
"LTX2VideoDownBlock3D",
),
"decoder_block_out_channels": (256, 512, 1024),
"layers_per_block": (4, 6, 6, 2, 2),
"decoder_layers_per_block": (5, 5, 5, 5),
"spatio_temporal_scaling": (True, True, True, True),
"decoder_spatio_temporal_scaling": (True, True, True),
"decoder_inject_noise": (False, False, False, False),
"downsample_type": (
"spatial",
"temporal",
"spatiotemporal",
"spatiotemporal",
),
"upsample_residual": (True, True, True),
"upsample_factor": (2, 2, 2),
"timestep_conditioning": False,
"patch_size": 4,
"patch_size_t": 1,
"resnet_norm_eps": 1e-6,
"encoder_causal": True,
"decoder_causal": False,
"encoder_spatial_padding_mode": "zeros",
"decoder_spatial_padding_mode": "reflect",
"spatial_compression_ratio": 32,
"temporal_compression_ratio": 8,
},
}
rename_dict = LTX_2_0_VIDEO_VAE_RENAME_DICT
special_keys_remap = LTX_2_0_VAE_SPECIAL_KEYS_REMAP
def get_ltx2_video_vae_config(
version: str = "2.0", timestep_conditioning: bool = False
) -> tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]:
if version == "test":
config = {
"model_id": "diffusers-internal-dev/dummy-ltx2",
"diffusers_config": {
"in_channels": 3,
"out_channels": 3,
"latent_channels": 128,
"block_out_channels": (256, 512, 1024, 2048),
"down_block_types": (
"LTX2VideoDownBlock3D",
"LTX2VideoDownBlock3D",
"LTX2VideoDownBlock3D",
"LTX2VideoDownBlock3D",
),
"decoder_block_out_channels": (256, 512, 1024),
"layers_per_block": (4, 6, 6, 2, 2),
"decoder_layers_per_block": (5, 5, 5, 5),
"spatio_temporal_scaling": (True, True, True, True),
"decoder_spatio_temporal_scaling": (True, True, True),
"decoder_inject_noise": (False, False, False, False),
"downsample_type": (
"spatial",
"temporal",
"spatiotemporal",
"spatiotemporal",
),
"upsample_residual": (True, True, True),
"upsample_factor": (2, 2, 2),
"timestep_conditioning": timestep_conditioning,
"patch_size": 4,
"patch_size_t": 1,
"resnet_norm_eps": 1e-6,
"encoder_causal": True,
"decoder_causal": False,
"encoder_spatial_padding_mode": "zeros",
"decoder_spatial_padding_mode": "reflect",
"spatial_compression_ratio": 32,
"temporal_compression_ratio": 8,
},
}
rename_dict = LTX_2_0_VIDEO_VAE_RENAME_DICT
special_keys_remap = LTX_2_0_VAE_SPECIAL_KEYS_REMAP
elif version == "2.0":
config = {
"model_id": "Lightricks/LTX-2",
"diffusers_config": {
"in_channels": 3,
"out_channels": 3,
"latent_channels": 128,
"block_out_channels": (256, 512, 1024, 2048),
"down_block_types": (
"LTX2VideoDownBlock3D",
"LTX2VideoDownBlock3D",
"LTX2VideoDownBlock3D",
"LTX2VideoDownBlock3D",
),
"decoder_block_out_channels": (256, 512, 1024),
"layers_per_block": (4, 6, 6, 2, 2),
"decoder_layers_per_block": (5, 5, 5, 5),
"spatio_temporal_scaling": (True, True, True, True),
"decoder_spatio_temporal_scaling": (True, True, True),
"decoder_inject_noise": (False, False, False, False),
"downsample_type": (
"spatial",
"temporal",
"spatiotemporal",
"spatiotemporal",
),
"upsample_type": ("spatiotemporal", "spatiotemporal", "spatiotemporal"),
"upsample_residual": (True, True, True),
"upsample_factor": (2, 2, 2),
"timestep_conditioning": timestep_conditioning,
"patch_size": 4,
"patch_size_t": 1,
"resnet_norm_eps": 1e-6,
"encoder_causal": True,
"decoder_causal": False,
"encoder_spatial_padding_mode": "zeros",
"decoder_spatial_padding_mode": "reflect",
"spatial_compression_ratio": 32,
"temporal_compression_ratio": 8,
},
}
rename_dict = LTX_2_0_VIDEO_VAE_RENAME_DICT
special_keys_remap = LTX_2_0_VAE_SPECIAL_KEYS_REMAP
elif version == "2.3":
config = {
"model_id": "Lightricks/LTX-2.3",
"diffusers_config": {
"in_channels": 3,
"out_channels": 3,
"latent_channels": 128,
"block_out_channels": (256, 512, 1024, 1024),
"down_block_types": (
"LTX2VideoDownBlock3D",
"LTX2VideoDownBlock3D",
"LTX2VideoDownBlock3D",
"LTX2VideoDownBlock3D",
),
"decoder_block_out_channels": (256, 512, 512, 1024),
"layers_per_block": (4, 6, 4, 2, 2),
"decoder_layers_per_block": (4, 6, 4, 2, 2),
"spatio_temporal_scaling": (True, True, True, True),
"decoder_spatio_temporal_scaling": (True, True, True, True),
"decoder_inject_noise": (False, False, False, False, False),
"downsample_type": (
"spatial",
"temporal",
"spatiotemporal",
"spatiotemporal",
),
"upsample_type": (
"spatiotemporal",
"spatiotemporal",
"temporal",
"spatial",
),
"upsample_residual": (False, False, False, False),
"upsample_factor": (2, 2, 1, 2),
"timestep_conditioning": timestep_conditioning,
"patch_size": 4,
"patch_size_t": 1,
"resnet_norm_eps": 1e-6,
"encoder_causal": True,
"decoder_causal": False,
"encoder_spatial_padding_mode": "zeros",
"decoder_spatial_padding_mode": "zeros",
"spatial_compression_ratio": 32,
"temporal_compression_ratio": 8,
},
}
rename_dict = LTX_2_3_VIDEO_VAE_RENAME_DICT
special_keys_remap = LTX_2_0_VAE_SPECIAL_KEYS_REMAP
return config, rename_dict, special_keys_remap
def convert_ltx2_video_vae(original_state_dict: Dict[str, Any]) -> Dict[str, Any]:
config, rename_dict, special_keys_remap = get_ltx2_video_vae_config()
def convert_ltx2_video_vae(
original_state_dict: Dict[str, Any],
version: str = "2.0",
timestep_conditioning: bool = False,
) -> Dict[str, Any]:
config, rename_dict, special_keys_remap = get_ltx2_video_vae_config(
version, timestep_conditioning
)
diffusers_config = config["diffusers_config"]
with init_empty_weights():
@@ -387,38 +689,66 @@ def convert_ltx2_video_vae(original_state_dict: Dict[str, Any]) -> Dict[str, Any
return vae
def get_ltx2_audio_vae_config() -> Tuple[
Dict[str, Any], Dict[str, Any], Dict[str, Any]
]:
config = {
"model_id": "diffusers-internal-dev/new-ltx-model",
"diffusers_config": {
"base_channels": 128,
"output_channels": 2,
"ch_mult": (1, 2, 4),
"num_res_blocks": 2,
"attn_resolutions": None,
"in_channels": 2,
"resolution": 256,
"latent_channels": 8,
"norm_type": "pixel",
"causality_axis": "height",
"dropout": 0.0,
"mid_block_add_attention": False,
"sample_rate": 16000,
"mel_hop_length": 160,
"is_causal": True,
"mel_bins": 64,
"double_z": True,
},
}
rename_dict = LTX_2_0_AUDIO_VAE_RENAME_DICT
special_keys_remap = LTX_2_0_AUDIO_VAE_SPECIAL_KEYS_REMAP
def get_ltx2_audio_vae_config(
version: str = "2.0",
) -> tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]:
if version == "2.0":
config = {
"model_id": "Lightricks/LTX-2",
"diffusers_config": {
"base_channels": 128,
"output_channels": 2,
"ch_mult": (1, 2, 4),
"num_res_blocks": 2,
"attn_resolutions": None,
"in_channels": 2,
"resolution": 256,
"latent_channels": 8,
"norm_type": "pixel",
"causality_axis": "height",
"dropout": 0.0,
"mid_block_add_attention": False,
"sample_rate": 16000,
"mel_hop_length": 160,
"is_causal": True,
"mel_bins": 64,
"double_z": True,
},
}
rename_dict = LTX_2_0_AUDIO_VAE_RENAME_DICT
special_keys_remap = LTX_2_0_AUDIO_VAE_SPECIAL_KEYS_REMAP
elif version == "2.3":
config = {
"model_id": "Lightricks/LTX-2.3",
"diffusers_config": {
"base_channels": 128,
"output_channels": 2,
"ch_mult": (1, 2, 4),
"num_res_blocks": 2,
"attn_resolutions": None,
"in_channels": 2,
"resolution": 256,
"latent_channels": 8,
"norm_type": "pixel",
"causality_axis": "height",
"dropout": 0.0,
"mid_block_add_attention": False,
"sample_rate": 16000,
"mel_hop_length": 160,
"is_causal": True,
"mel_bins": 64,
"double_z": True,
}, # Same config as LTX-2.0
}
rename_dict = LTX_2_0_AUDIO_VAE_RENAME_DICT
special_keys_remap = LTX_2_0_AUDIO_VAE_SPECIAL_KEYS_REMAP
return config, rename_dict, special_keys_remap
def convert_ltx2_audio_vae(original_state_dict: Dict[str, Any]) -> Dict[str, Any]:
config, rename_dict, special_keys_remap = get_ltx2_audio_vae_config()
def convert_ltx2_audio_vae(
original_state_dict: Dict[str, Any], version: str = "2.0"
) -> Dict[str, Any]:
config, rename_dict, special_keys_remap = get_ltx2_audio_vae_config(version)
diffusers_config = config["diffusers_config"]
with init_empty_weights():
@@ -443,32 +773,87 @@ def convert_ltx2_audio_vae(original_state_dict: Dict[str, Any]) -> Dict[str, Any
return vae
def get_ltx2_vocoder_config() -> Tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]:
config = {
"model_id": "diffusers-internal-dev/new-ltx-model",
"diffusers_config": {
"in_channels": 128,
"hidden_channels": 1024,
"out_channels": 2,
"upsample_kernel_sizes": [16, 15, 8, 4, 4],
"upsample_factors": [6, 5, 2, 2, 2],
"resnet_kernel_sizes": [3, 7, 11],
"resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
"leaky_relu_negative_slope": 0.1,
"output_sampling_rate": 24000,
},
}
rename_dict = LTX_2_0_VOCODER_RENAME_DICT
special_keys_remap = LTX_2_0_VOCODER_SPECIAL_KEYS_REMAP
def get_ltx2_vocoder_config(
version: str = "2.0",
) -> tuple[Dict[str, Any], dict[str, Any], dict[str, Any]]:
if version == "2.0":
config = {
"model_id": "Lightricks/LTX-2",
"diffusers_config": {
"in_channels": 128,
"hidden_channels": 1024,
"out_channels": 2,
"upsample_kernel_sizes": [16, 15, 8, 4, 4],
"upsample_factors": [6, 5, 2, 2, 2],
"resnet_kernel_sizes": [3, 7, 11],
"resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
"act_fn": "leaky_relu",
"leaky_relu_negative_slope": 0.1,
"antialias": False,
"final_act_fn": "tanh",
"final_bias": True,
"output_sampling_rate": 24000,
},
}
rename_dict = LTX_2_0_VOCODER_RENAME_DICT
special_keys_remap = LTX_2_0_VOCODER_SPECIAL_KEYS_REMAP
elif version == "2.3":
config = {
"model_id": "Lightricks/LTX-2.3",
"diffusers_config": {
"in_channels": 128,
"hidden_channels": 1536,
"out_channels": 2,
"upsample_kernel_sizes": [11, 4, 4, 4, 4, 4],
"upsample_factors": [5, 2, 2, 2, 2, 2],
"resnet_kernel_sizes": [3, 7, 11],
"resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
"act_fn": "snakebeta",
"leaky_relu_negative_slope": 0.1,
"antialias": True,
"antialias_ratio": 2,
"antialias_kernel_size": 12,
"final_act_fn": None,
"final_bias": False,
"bwe_in_channels": 128,
"bwe_hidden_channels": 512,
"bwe_out_channels": 2,
"bwe_upsample_kernel_sizes": [12, 11, 4, 4, 4],
"bwe_upsample_factors": [6, 5, 2, 2, 2],
"bwe_resnet_kernel_sizes": [3, 7, 11],
"bwe_resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
"bwe_act_fn": "snakebeta",
"bwe_leaky_relu_negative_slope": 0.1,
"bwe_antialias": True,
"bwe_antialias_ratio": 2,
"bwe_antialias_kernel_size": 12,
"bwe_final_act_fn": None,
"bwe_final_bias": False,
"filter_length": 512,
"hop_length": 80,
"window_length": 512,
"num_mel_channels": 64,
"input_sampling_rate": 16000,
"output_sampling_rate": 48000,
},
}
rename_dict = LTX_2_3_VOCODER_RENAME_DICT
special_keys_remap = LTX_2_3_VOCODER_SPECIAL_KEYS_REMAP
return config, rename_dict, special_keys_remap
def convert_ltx2_vocoder(original_state_dict: Dict[str, Any]) -> Dict[str, Any]:
config, rename_dict, special_keys_remap = get_ltx2_vocoder_config()
def convert_ltx2_vocoder(
original_state_dict: Dict[str, Any], version: str
) -> Dict[str, Any]:
config, rename_dict, special_keys_remap = get_ltx2_vocoder_config(version)
diffusers_config = config["diffusers_config"]
if version == "2.3":
vocoder_cls = LTX2VocoderWithBWE
else:
vocoder_cls = LTX2Vocoder
with init_empty_weights():
vocoder = LTX2Vocoder.from_config(diffusers_config)
vocoder = vocoder_cls.from_config(diffusers_config)
# Handle official code --> diffusers key remapping via the remap dict
for key in list(original_state_dict.keys()):
@@ -499,13 +884,17 @@ def get_model_state_dict_from_combined_ckpt(
model_state_dict = {}
for param_name, param in combined_ckpt.items():
if param_name.startswith(prefix):
model_state_dict[param_name.replace(prefix, "")] = param
model_state_dict[param_name.removeprefix(prefix)] = param
if prefix == "model.diffusion_model.":
# Some checkpoints store the text connector projection outside the diffusion model prefix.
connector_key = "text_embedding_projection.aggregate_embed.weight"
if connector_key in combined_ckpt and connector_key not in model_state_dict:
model_state_dict[connector_key] = combined_ckpt[connector_key]
connector_prefixes = ["text_embedding_projection"]
for param_name, param in combined_ckpt.items():
for prefix in connector_prefixes:
if param_name.startswith(prefix):
# Check to make sure we're not overwriting an existing key
if param_name not in model_state_dict:
model_state_dict[param_name] = combined_ckpt[param_name]
return model_state_dict
@@ -586,9 +975,12 @@ def convert_comfy_gemma3_to_transformers(sd: dict):
def convert_lora_original_to_diffusers(
lora_state_dict: Dict[str, Any],
version: str = "2.0",
) -> Dict[str, Any]:
out: Dict[str, Any] = {}
rename_dict = LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT
if version == "2.3":
rename_dict = LTX_2_3_TRANSFORMER_KEYS_RENAME_DICT
for k, v in lora_state_dict.items():
# Keep the "diffusion_model." prefix as-is, but apply the transformer remaps to the rest
@@ -617,10 +1009,14 @@ def convert_lora_original_to_diffusers(
def convert_lora_diffusers_to_original(
lora_state_dict: Dict[str, Any],
version: str = "2.0",
) -> Dict[str, Any]:
out: Dict[str, Any] = {}
rename_dict = LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT
if version == "2.3":
rename_dict = LTX_2_3_TRANSFORMER_KEYS_RENAME_DICT
inv_rename = {v: k for k, v in LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT.items()}
inv_rename = {v: k for k, v in rename_dict.items()}
inv_items = sorted(inv_rename.items(), key=lambda kv: len(kv[0]), reverse=True)
for k, v in lora_state_dict.items():

127
extensions_built_in/diffusion_models/ltx2/ltx2.py
View File

@@ -21,6 +21,7 @@ from toolkit.util.quantize import quantize, get_qtype, quantize_model
from toolkit.memory_management import MemoryManager
from safetensors.torch import load_file
from PIL import Image
import huggingface_hub
try:
from diffusers import LTX2Pipeline, LTX2ImageToVideoPipeline
@@ -34,7 +35,7 @@ try:
Gemma3ForConditionalGeneration,
GemmaTokenizerFast,
)
from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder, LTX2VocoderWithBWE
from diffusers.pipelines.ltx2.connectors import LTX2TextConnectors
from .convert_ltx2_to_diffusers import (
get_model_state_dict_from_combined_ckpt,
@@ -76,6 +77,9 @@ dit_prefix = "model.diffusion_model."
vae_prefix = "vae."
audio_vae_prefix = "audio_vae."
vocoder_prefix = "vocoder."
base_te_path = "google/gemma-3-12b-it-qat-q4_0-unquantized"
HF_TOKEN = os.getenv("HF_TOKEN", None)
def new_save_image_function(
@@ -192,6 +196,8 @@ class AudioProcessor(torch.nn.Module):
class LTX2Model(BaseModel):
arch = "ltx2"
ltx_version = "2.0"
ltx_te_path = None
def __init__(
self,
@@ -213,7 +219,7 @@ class LTX2Model(BaseModel):
# use the new format on this new model by default
self.use_old_lokr_format = False
self.audio_processor = None
# gemma needs left side padding
self.te_padding_side = "left"
@@ -234,6 +240,21 @@ class LTX2Model(BaseModel):
combined_state_dict = None
self.print_and_status_update("Loading transformer")
if not os.path.exists(model_path) and model_path.endswith(".safetensors"):
# download the model from the Hugging Face Hub if it is not a local path
splits = model_path.split("/")
if len(splits) != 3:
raise ValueError(
f"Invalid model path: {model_path}. Must be in the format 'repo_id/repo/filename.safetensors' to download from the Hugging Face Hub."
)
# download the model from the hub
model_path = huggingface_hub.hf_hub_download(
repo_id="/".join(splits[:2]),
filename=splits[2],
token=HF_TOKEN,
)
# if we have a safetensors file it is a mono checkpoint
if os.path.exists(model_path) and model_path.endswith(".safetensors"):
combined_state_dict = load_file(model_path)
@@ -243,7 +264,7 @@ class LTX2Model(BaseModel):
original_dit_ckpt = get_model_state_dict_from_combined_ckpt(
combined_state_dict, dit_prefix
)
transformer = convert_ltx2_transformer(original_dit_ckpt)
transformer = convert_ltx2_transformer(original_dit_ckpt, version=self.ltx_version)
transformer = transformer.to(dtype)
else:
transformer_path = model_path
@@ -297,9 +318,7 @@ class LTX2Model(BaseModel):
and self.model_config.te_name_or_path.endswith(".safetensors")
):
# load from comfyui gemma3 checkpoint
tokenizer = GemmaTokenizerFast.from_pretrained(
"Lightricks/LTX-2", subfolder="tokenizer"
)
tokenizer = GemmaTokenizerFast.from_pretrained(base_te_path)
with init_empty_weights():
text_encoder = Gemma3ForConditionalGeneration(
@@ -371,15 +390,21 @@ class LTX2Model(BaseModel):
del te_state_dict
flush()
else:
te_subfolder = "tokenizer"
if self.model_config.te_name_or_path is not None:
te_path = self.model_config.te_name_or_path
else:
te_path = base_model_path
if self.ltx_te_path is not None:
# pull from original repo for 2.3
te_path = self.ltx_te_path
te_subfolder = None
else:
te_path = base_te_path
tokenizer = GemmaTokenizerFast.from_pretrained(
te_path, subfolder="tokenizer"
te_path, subfolder=te_subfolder
)
text_encoder = Gemma3ForConditionalGeneration.from_pretrained(
te_path, subfolder="text_encoder", dtype=dtype
te_path, subfolder=te_subfolder, dtype=dtype
)
# remove the vision tower
@@ -413,22 +438,22 @@ class LTX2Model(BaseModel):
original_vae_ckpt = get_model_state_dict_from_combined_ckpt(
combined_state_dict, vae_prefix
)
vae = convert_ltx2_video_vae(original_vae_ckpt).to(dtype)
vae = convert_ltx2_video_vae(original_vae_ckpt, version=self.ltx_version).to(dtype)
del original_vae_ckpt
original_audio_vae_ckpt = get_model_state_dict_from_combined_ckpt(
combined_state_dict, audio_vae_prefix
)
audio_vae = convert_ltx2_audio_vae(original_audio_vae_ckpt).to(dtype)
audio_vae = convert_ltx2_audio_vae(original_audio_vae_ckpt, version=self.ltx_version).to(dtype)
del original_audio_vae_ckpt
original_connectors_ckpt = get_model_state_dict_from_combined_ckpt(
combined_state_dict, dit_prefix
)
connectors = convert_ltx2_connectors(original_connectors_ckpt).to(dtype)
connectors = convert_ltx2_connectors(original_connectors_ckpt, version=self.ltx_version).to(dtype)
del original_connectors_ckpt
original_vocoder_ckpt = get_model_state_dict_from_combined_ckpt(
combined_state_dict, vocoder_prefix
)
vocoder = convert_ltx2_vocoder(original_vocoder_ckpt).to(dtype)
vocoder = convert_ltx2_vocoder(original_vocoder_ckpt, version=self.ltx_version).to(dtype)
del original_vocoder_ckpt
del combined_state_dict
flush()
@@ -444,7 +469,11 @@ class LTX2Model(BaseModel):
base_model_path, subfolder="connectors", torch_dtype=dtype
)
vocoder = LTX2Vocoder.from_pretrained(
vocoder_cls = LTX2Vocoder
if self.ltx_version == "2.3":
vocoder_cls = LTX2VocoderWithBWE
vocoder = vocoder_cls.from_pretrained(
base_model_path, subfolder="vocoder", torch_dtype=dtype
)
@@ -630,10 +659,21 @@ class LTX2Model(BaseModel):
tile_sample_stride_width=224,
tile_sample_stride_num_frames=4,
)
# We only encode and store the minimum prompt tokens, but need them padded to 1024 for LTX2
conditional_embeds = self.pad_embeds(conditional_embeds)
unconditional_embeds = self.pad_embeds(unconditional_embeds)
if self.ltx_version == "2.3":
extra['stg_scale'] = 1.0
extra['modality_scale'] = 3.0
extra['guidance_rescale'] = 0.7
extra['audio_guidance_scale'] = 7.0
extra['audio_stg_scale'] = 1.0
extra['audio_modality_scale'] = 3.0
extra['audio_guidance_rescale'] = 0.7
extra['spatio_temporal_guidance_blocks'] = [28]
extra['use_cross_timestep'] = True # they dont set this in some examples in diffusers, but I believe it should always be true for 2.3
video, audio = pipeline(
prompt_embeds=conditional_embeds.text_embeds.to(
@@ -738,7 +778,7 @@ class LTX2Model(BaseModel):
latents_std = self.pipeline.audio_vae.latents_std
output_tensor = (output_tensor - latents_mean) / latents_std
return output_tensor
def pad_embeds(self, embeds: PromptEmbeds):
# ltx-2 connector requires 1024 tokens for good results. Any smaller and it degrades.
target_length = 1024
@@ -773,7 +813,7 @@ class LTX2Model(BaseModel):
with torch.no_grad():
if self.model.device == torch.device("cpu"):
self.model.to(self.device_torch)
# We only encode and store the minimum prompt tokens, but need them padded to 1024 for LTX2
text_embeddings = self.pad_embeds(text_embeddings)
@@ -884,16 +924,17 @@ class LTX2Model(BaseModel):
if self.pipeline.connectors.device != self.transformer.device:
self.pipeline.connectors.to(self.transformer.device)
# TODO this is how diffusers does this on inference, not sure I understand why, check this
additive_attention_mask = (
1 - text_embeddings.attention_mask.to(self.transformer.dtype)
) * -1000000.0
tokenizer_padding_side = "left" # Padding side for default Gemma3-12B text encoder
if getattr(self, "tokenizer", None) is not None:
tokenizer_padding_side = getattr(self.tokenizer, "padding_side", "left")
(
connector_prompt_embeds,
connector_audio_prompt_embeds,
connector_attention_mask,
) = self.pipeline.connectors(
text_embeddings.text_embeds, additive_attention_mask, additive_mask=True
text_embeddings.text_embeds,
text_embeddings.attention_mask.to(self.transformer.dtype),
padding_side=tokenizer_padding_side,
)
# compute video and audio positional ids
@@ -908,6 +949,11 @@ class LTX2Model(BaseModel):
audio_coords = self.transformer.audio_rope.prepare_audio_coords(
audio_latents.shape[0], audio_num_frames, audio_latents.device
)
# use_cross_timestep - Whether to use the cross modality (audio is the cross modality of video, and vice versa) sigma when
# calculating the cross attention modulation parameters. `True` is the newer (e.g. LTX-2.3) behavior;
# `False` is the legacy LTX-2.0 behavior.
use_cross_timestep = self.ltx_version == "2.3"
noise_pred_video, noise_pred_audio = self.transformer(
hidden_states=packed_latents,
@@ -915,6 +961,7 @@ class LTX2Model(BaseModel):
encoder_hidden_states=connector_prompt_embeds,
audio_encoder_hidden_states=connector_audio_prompt_embeds,
timestep=video_timestep,
sigma=video_timestep, # Used by LTX-2.3
audio_timestep=timestep,
encoder_attention_mask=connector_attention_mask,
audio_encoder_attention_mask=connector_attention_mask,
@@ -925,7 +972,10 @@ class LTX2Model(BaseModel):
audio_num_frames=audio_num_frames,
video_coords=video_coords,
audio_coords=audio_coords,
# rope_interpolation_scale=rope_interpolation_scale,
isolate_modalities=False,
spatio_temporal_guidance_blocks=None,
perturbation_mask=None,
use_cross_timestep=use_cross_timestep,
attention_kwargs=None,
return_dict=False,
)
@@ -969,7 +1019,7 @@ class LTX2Model(BaseModel):
)
text_input_ids = text_inputs.input_ids
prompt_attention_mask = text_inputs.attention_mask
text_input_ids = text_input_ids.to(device)
prompt_attention_mask = prompt_attention_mask.to(device)
@@ -980,27 +1030,18 @@ class LTX2Model(BaseModel):
)
text_encoder_hidden_states = text_encoder_outputs.hidden_states
text_encoder_hidden_states = torch.stack(text_encoder_hidden_states, dim=-1)
sequence_lengths = prompt_attention_mask.sum(dim=-1)
prompt_embeds = self.pipeline._pack_text_embeds(
text_encoder_hidden_states,
sequence_lengths,
device=device,
padding_side=self.tokenizer[0].padding_side,
scale_factor=scale_factor,
)
prompt_embeds = prompt_embeds.to(dtype=self.torch_dtype)
prompt_embeds = text_encoder_hidden_states.flatten(2, 3).to(
dtype=self.torch_dtype
) # Pack to 3D
# duplicate text embeddings for each generation per prompt, using mps friendly method
_, seq_len, _ = prompt_embeds.shape
prompt_embeds = prompt_embeds.repeat(1, 1, 1)
prompt_embeds = prompt_embeds.view(
batch_size * 1, seq_len, -1
)
prompt_embeds = prompt_embeds.view(batch_size * 1, seq_len, -1)
prompt_attention_mask = prompt_attention_mask.view(batch_size, -1)
prompt_attention_mask = prompt_attention_mask.repeat(1, 1)
pe = PromptEmbeds([prompt_embeds, None])
pe.attention_mask = prompt_attention_mask
return pe
@@ -1038,13 +1079,19 @@ class LTX2Model(BaseModel):
for key, value in state_dict.items():
new_key = key.replace("transformer.", "diffusion_model.")
new_sd[new_key] = value
new_sd = convert_lora_diffusers_to_original(new_sd)
new_sd = convert_lora_diffusers_to_original(new_sd, version=self.ltx_version)
return new_sd
def convert_lora_weights_before_load(self, state_dict):
state_dict = convert_lora_original_to_diffusers(state_dict)
state_dict = convert_lora_original_to_diffusers(state_dict, version=self.ltx_version)
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("diffusion_model.", "transformer.")
new_sd[new_key] = value
return new_sd
class LTX23Model(LTX2Model):
arch = "ltx2.3"
ltx_version = "2.3"
ltx_te_path = base_te_path

2
requirements.txt
View File

@@ -1,6 +1,6 @@
torchao==0.10.0
safetensors
git+https://github.com/huggingface/diffusers@8600b4c10d67b0ce200f664204358747bd53c775
git+https://github.com/huggingface/diffusers@072d15ee4289ffdc3aa9d65f8b94bc9271319d21
transformers==4.57.3
lycoris-lora==1.8.3
flatten_json

27
ui/src/app/jobs/new/options.ts
View File

@@ -668,6 +668,33 @@ export const modelArchs: ModelArch[] = [
disableSections: ['network.conv'],
additionalSections: ['sample.ctrl_img', 'datasets.num_frames', 'model.layer_offloading', 'model.low_vram', 'datasets.do_audio', 'datasets.audio_normalize', 'datasets.audio_preserve_pitch', 'datasets.do_i2v', 'train.audio_loss_multiplier'],
},
{
name: 'ltx2.3',
label: 'LTX-2.3',
group: 'video',
isVideoModel: true,
defaults: {
// default updates when [selected, unselected] in the UI
'config.process[0].model.name_or_path': ['Lightricks/LTX-2.3/ltx-2.3-22b-dev.safetensors', defaultNameOrPath],
'config.process[0].model.quantize': [true, false],
'config.process[0].model.quantize_te': [true, false],
'config.process[0].model.low_vram': [true, false],
'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
'config.process[0].sample.num_frames': [121, 1],
'config.process[0].sample.fps': [24, 1],
'config.process[0].sample.width': [768, 1024],
'config.process[0].sample.height': [768, 1024],
'config.process[0].train.audio_loss_multiplier': [1.0, undefined],
'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
'config.process[0].datasets[x].cache_latents_to_disk': [true, false],
'config.process[0].datasets[x].do_i2v': [false, undefined],
'config.process[0].datasets[x].do_audio': [true, undefined],
'config.process[0].datasets[x].fps': [24, undefined],
},
disableSections: ['network.conv'],
additionalSections: ['sample.ctrl_img', 'datasets.num_frames', 'model.layer_offloading', 'model.low_vram', 'datasets.do_audio', 'datasets.audio_normalize', 'datasets.audio_preserve_pitch', 'datasets.do_i2v', 'train.audio_loss_multiplier'],
},
{
name: 'flux2_klein_4b',
label: 'FLUX.2-klein-base-4B',

2
version.py
View File

@@ -1 +1 @@
VERSION = "0.7.24"
VERSION = "0.7.25"