From 859635e95bd6af68ddda638870d278614a4ffea4 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sun, 22 Mar 2026 17:56:59 -0600 Subject: [PATCH] Add support for training LTX 2.3 (#745) * 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 --- dgx_requirements.txt | 2 +- .../diffusion_models/__init__.py | 3 +- .../diffusion_models/ltx2/__init__.py | 2 +- .../ltx2/convert_ltx2_to_diffusers.py | 778 +++++++++++++----- .../diffusion_models/ltx2/ltx2.py | 127 ++- requirements.txt | 2 +- ui/src/app/jobs/new/options.ts | 27 + version.py | 2 +- 8 files changed, 707 insertions(+), 236 deletions(-) diff --git a/dgx_requirements.txt b/dgx_requirements.txt index d540b79..1906121 100644 --- a/dgx_requirements.txt +++ b/dgx_requirements.txt @@ -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 diff --git a/extensions_built_in/diffusion_models/__init__.py b/extensions_built_in/diffusion_models/__init__.py index 88d1bed..ed4ffd7 100644 --- a/extensions_built_in/diffusion_models/__init__.py +++ b/extensions_built_in/diffusion_models/__init__.py @@ -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, diff --git a/extensions_built_in/diffusion_models/ltx2/__init__.py b/extensions_built_in/diffusion_models/ltx2/__init__.py index 6ac6737..44b8407 100644 --- a/extensions_built_in/diffusion_models/ltx2/__init__.py +++ b/extensions_built_in/diffusion_models/ltx2/__init__.py @@ -1 +1 @@ -from .ltx2 import LTX2Model \ No newline at end of file +from .ltx2 import LTX2Model, LTX23Model \ No newline at end of file diff --git a/extensions_built_in/diffusion_models/ltx2/convert_ltx2_to_diffusers.py b/extensions_built_in/diffusion_models/ltx2/convert_ltx2_to_diffusers.py index 7ee8af2..2c6447b 100644 --- a/extensions_built_in/diffusion_models/ltx2/convert_ltx2_to_diffusers.py +++ b/extensions_built_in/diffusion_models/ltx2/convert_ltx2_to_diffusers.py @@ -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(): diff --git a/extensions_built_in/diffusion_models/ltx2/ltx2.py b/extensions_built_in/diffusion_models/ltx2/ltx2.py index a6d1e6d..e19cc3e 100644 --- a/extensions_built_in/diffusion_models/ltx2/ltx2.py +++ b/extensions_built_in/diffusion_models/ltx2/ltx2.py @@ -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 diff --git a/requirements.txt b/requirements.txt index 2ab3621..54f05f2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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 diff --git a/ui/src/app/jobs/new/options.ts b/ui/src/app/jobs/new/options.ts index e4001cc..bf88725 100644 --- a/ui/src/app/jobs/new/options.ts +++ b/ui/src/app/jobs/new/options.ts @@ -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', diff --git a/version.py b/version.py index 9d39d73..ff86031 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.7.24" +VERSION = "0.7.25"