AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-23 09:09:01 +03:00
Code Activity
Files
linked-issue-self-opened
diffusers/scripts/convert_ltx2_to_diffusers.py
Alex Artyomov 82e309a2b1 Add LTX-2.5 DFR pipeline (keyframe slots, spatial detailing, tiled temporal rounds) (#14567) 
* Add LTX-2.5 DFR pipeline with generated keyframe slots

Ports DFRPipeline from the Lightricks reference. Stage 1 generates video plus
single-pixel-frame keyframe slots at a fraction of the requested resolution on
a VAE-aligned segment grid; both are spatially latent-upsampled and stage 2
re-denoises at twice that resolution with the slots re-attached and an optional
spatial detailing IC-LoRA active for that stage only. Optional temporal x2/x4
refine rounds tile the canvas at shared keyframes and densify with ancestral
Euler. With spatial_upscalings=2 a full-resolution detailing epilogue follows
the rounds.

The transformer already stored keyframes_abs_pos_embedding for load/save; this
wires it into the forward through a new video_keyframes_mask argument, which
only a keyframes-aware pipeline passes, so other pipelines are unaffected.

The epilogue denoises the whole canvas in one loop and tiles the transformer
call inside it, so every Euler step steps a canvas whose tiles have already
agreed on their overlaps. Spatial tiles blend under a trapezoidal mask, since
neither side of a height or width border holds a known answer. Temporal tiles
are cut on the keyframe seams the last refine round stitched on: both windows
reproduce a shared keyframe there, so the later one drops its run-up under a
rectangular mask rather than averaging it. Conditionings are attached once on
the whole canvas and filtered per tile at the token level, and a keyframe two
windows share is one token they both read.

The epilogue is handed its keyframes rather than asked to generate them. Each
carry plane is decoded on its own -- the VAE is causal, so a stacked decode
would bleed neighbours -- then Lanczos-stretched x2 in RGB and encoded again at
the output resolution, and pinned fully clean. Only the video latent is
spatially upsampled.

Conditioning fps is snapped to 60 above 30 rather than merely capped there, at
every stage. RoPE time is pixel_frame / fps, and the transformer is trained
around 24/25/30 and 60; a temporal round taking 24 fps to 48 lands between
those, and it shows as stutter at the latent borders. Playback fps is
unchanged, so 24 fps with one round still ships 48 fps.

A condition's index is read on the canvas num_frames asks for, and the moment
it names is carried onto each refine round's longer canvas by scaling its pixel
position by 2**round. The scaled position does not generally land on a latent
boundary, so it travels as a pixel index rather than through the public latent
index; a keyframe conditioning is appended as extra tokens instead of being
spliced into the base grid, so it does not need to.

height and width must be divisible by 2**spatial_upscalings times the VAE's
spatial compression ratio, which makes 4K 3840x2176 rather than 3840x2160. That
rule is checked ahead of the looser one every LTX-2 pipeline applies, so the
error names the divisor a DFR caller actually has to satisfy.

Four details are easy to get wrong, and each is covered by a test after showing
up as a visible seam at a tile handover:

- The ancestral step injects noise into every token, so the conditioning blend
  has to be re-applied afterwards. Skipping it lets the strength-0.95 anchor
  keyframes erode over the schedule, and those anchors are the only thing
  pinning adjacent tiles onto the same content.
- Velocity is converted to x0 with each token's own noise level, not the scalar
  schedule sigma: a token held at strength s sits at (1 - s) * sigma.
- Each tile draws its ancestral noise from a generator seeded
  seed + 1000 * round + tile, kept separate from the main generator so the
  draws do not consume state the next tile's initial noising reads.
- Two tiles invent the slot that falls in the later one's dropped lead-in. The
  stitch keeps the earlier tile's frames there, so the earlier tile's copy is
  the one the canvas holds, and the one the next round must anchor on.

* Split the DFR pipeline into composable stages

Addresses the review on #14567: use the same compose pattern as the other LTX
two-stage pipelines, get ancestral Euler from an existing scheduler, and put the
temporal rounds in their own pipeline so the schedule is not switched mid-call.

`LTX2DFRPipeline.__call__` is now one denoise pass at `height` x `width`. Callers
compose stage 1, `LTX2LatentUpsamplePipeline`, stage 2 and each temporal round,
and a documented recipe is the copy-paste 1080p path. The recipe knobs
(`spatial_upscalings`, `temporal_upscalings`, `detailing_lora_adapter_name`) and
the required upsampler components are gone; `height`/`width` are this pass, not
the final output.

`ancestral_euler_step` is replaced by `LTXEulerAncestralRFScheduler.step` plus a
re-application of the conditioning blend, which ancestral noise would otherwise
erode on the strength-0.95 seam anchors. The new `LTX2DFRTemporalRefinePipeline`
owns that scheduler and one round; stage 1, stage 2 and the epilogue stay on
`FlowMatchEulerDiscreteScheduler`. It refuses any other scheduler rather than
silently taking a deterministic step and returning a softer canvas.

Pack/unpack, `prepare_latents`, `denoise` and `encode_conditions` move to
`LTX2DFRCoreMixin` so neither pipeline subclasses the other. Public latents are
raw on both sides of every boundary, `output_type="latent"` returns the untrimmed
canvas so a slot on the pad is not dropped, and `trim_canvas` does the trim
before decode. `LTX2DFRPipelineOutput` carries `keyframes` and
`keyframe_positions`, which cannot be re-derived after a round.

Verified against the pre-split implementation: bit-exact on dummy components in
fp32 and bf16, with and without the IC-LoRA reference, over one and two rounds;
and within one bf16 ulp on the real checkpoint, where the only difference is that
the split normalizes upsampled latents in fp32 rather than bf16.

Also in this pass:

- `__call__` takes `video_tiles` (the `epilogue_tiles` layout) instead of a
  resolved token plan. Resolving one needs the RoPE coordinates that only exist
  once `prepare_latents` has run, so a caller could not build the plan at all.
- `rebuild_epilogue_keyframes` is public and takes and returns raw latents. The
  composed epilogue needs it, so it was public API in practice while named
  private, and its normalized return forced callers into `_denormalize_latents`.
- Drop prompt enhancement and `num_videos_per_prompt` from the temporal pipeline.
  Enhancement belongs to stage 1 -- re-running it would denoise the canvas under
  a different prompt than the one that generated it -- and the batch is set by
  the incoming latent canvas, so `num_videos_per_prompt > 1` only ever raised.
- Fix the docs recipe: `requested_frames` counted latent frames where
  `trim_canvas` wants pixel frames, truncating a 241-frame render to 25; the
  three pipelines share components, so place them together instead of offloading
  one and leaving `temporal_latent_upsampler` off the device; and the detailing
  IC-LoRA is applied at 0.5, the strength the reference hardcodes.

- Pass `crf=0` on the conditions in the refine-round test. The default CRF sends
  the image through H.264 re-compression, which needs PyAV, so the test failed on
  any environment without it while testing nothing about re-compression.

* Drop LTX2DFRCoreMixin in favor of # Copied from

Each DFR pipeline now owns pack/unpack, prepare_latents, and denoise so the
classes stay self-contained. Shared helpers copy from LTX2Pipeline /
LTX2ConditionPipeline, and DFR-specific methods copy from LTX2DFRPipeline.
dfr_core.py keeps only the constants and canvas helpers.

* inline _finalize_output

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* remove dfr_core.py

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* make return_dict=False return a tuple of 4, consistent with diffusers convention; update docs and pipeline output accordingly

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* remove the dfr_layout test suite

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

---------

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
Co-authored-by: YiYi Xu <yixu310@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
2026-08-30 18:23:12 -10:00

1704 lines
72 KiB
Python
Raw Permalink Blame History

import argparse
import os
from contextlib import nullcontext
from typing import Any
import safetensors.torch
import torch
from accelerate import init_empty_weights
from huggingface_hub import hf_hub_download
from transformers import AutoConfig, AutoModelForImageTextToText, AutoProcessor, AutoTokenizer
from diffusers import (
AutoencoderKLLTX2Audio,
AutoencoderKLLTX2Video,
FlowMatchEulerDiscreteScheduler,
LTX2LatentUpsamplePipeline,
LTX2Pipeline,
LTX2VideoDiffusionDecoderModel,
LTX2VideoTransformer3DModel,
)
from diffusers.pipelines.ltx2 import (
LTX2DurationHead,
LTX2LatentUpsamplerModel,
LTX2TextConnectors,
LTX2Vocoder,
LTX2VocoderWithBWE,
)
from diffusers.utils.import_utils import is_accelerate_available
CTX = init_empty_weights if is_accelerate_available() else nullcontext
LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT = {
# Input Patchify Projections
"patchify_proj": "proj_in",
"audio_patchify_proj": "audio_proj_in",
# Modulation Parameters
# Handle adaln_single --> time_embed, audioln_single --> audio_time_embed separately as the original keys are
# substrings of the other modulation parameters below
"av_ca_video_scale_shift_adaln_single": "av_cross_attn_video_scale_shift",
"av_ca_a2v_gate_adaln_single": "av_cross_attn_video_a2v_gate",
"av_ca_audio_scale_shift_adaln_single": "av_cross_attn_audio_scale_shift",
"av_ca_v2a_gate_adaln_single": "av_cross_attn_audio_v2a_gate",
# Transformer Blocks
# Per-Block Cross Attention Modulatin Parameters
"scale_shift_table_a2v_ca_video": "video_a2v_cross_attn_scale_shift_table",
"scale_shift_table_a2v_ca_audio": "audio_a2v_cross_attn_scale_shift_table",
# Attention QK Norms
"q_norm": "norm_q",
"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",
"down_blocks.1": "down_blocks.0.downsamplers.0",
"down_blocks.2": "down_blocks.1",
"down_blocks.3": "down_blocks.1.downsamplers.0",
"down_blocks.4": "down_blocks.2",
"down_blocks.5": "down_blocks.2.downsamplers.0",
"down_blocks.6": "down_blocks.3",
"down_blocks.7": "down_blocks.3.downsamplers.0",
"down_blocks.8": "mid_block",
# Decoder
"up_blocks.0": "mid_block",
"up_blocks.1": "up_blocks.0.upsamplers.0",
"up_blocks.2": "up_blocks.0",
"up_blocks.3": "up_blocks.1.upsamplers.0",
"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",
"per_channel_statistics.mean-of-means": "latents_mean",
"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.5's diffusion decoder replaces the conv decoder while keeping the same encoder, so only the
# `decoder.*` half of the VAE checkpoint is renamed with these rules.
LTX_2_5_DIFFUSION_DECODER_RENAME_DICT = {
# The original `t_embedder` *is* diffusers' `PixArtAlphaCombinedTimestepSizeEmbeddings`, saved under
# shorter names, so only its two Linears need renaming.
"t_embedder.mlp.0.": "t_embedder.timestep_embedder.linear_1.",
"t_embedder.mlp.2.": "t_embedder.timestep_embedder.linear_2.",
".attn.proj.": ".attn.to_out.0.",
".attn.q_norm.": ".attn.norm_q.",
".attn.k_norm.": ".attn.norm_k.",
}
# Where a checkpoint carries static AdaLN gates, each is folded into the Linear it gates (W <- g * W) and
# dropped, because the decoder's residuals are ungated. Maps a renamed parameter to its gate's suffix.
LTX_2_5_DIFFUSION_DECODER_GATE_FOLD_TARGETS = {
".attn.to_out.0.weight": ".gate_msa",
".attn.to_out.0.bias": ".gate_msa",
".mlp.w_down.weight": ".gate_mlp",
".context_proj.weight": ".gate_ctx",
".context_proj.bias": ".gate_ctx",
}
LTX_2_0_AUDIO_VAE_RENAME_DICT = {
"per_channel_statistics.mean-of-means": "latents_mean",
"per_channel_statistics.std-of-means": "latents_std",
}
LTX_2_0_VOCODER_RENAME_DICT = {
"ups": "upsamplers",
"resblocks": "resnets",
"conv_pre": "conv_in",
"conv_post": "conv_out",
}
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",
}
def update_state_dict_inplace(state_dict: dict[str, Any], old_key: str, new_key: str) -> None:
state_dict[new_key] = state_dict.pop(old_key)
def remove_keys_inplace(key: str, state_dict: dict[str, Any]) -> None:
state_dict.pop(key)
def convert_ltx2_transformer_adaln_single(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 key.startswith("adaln_single."):
new_key = key.replace("adaln_single.", "time_embed.")
param = state_dict.pop(key)
state_dict[new_key] = param
if key.startswith("audio_adaln_single."):
new_key = key.replace("audio_adaln_single.", "audio_time_embed.")
param = state_dict.pop(key)
state_dict[new_key] = param
return
def convert_ltx2_audio_vae_per_channel_statistics(key: str, state_dict: dict[str, Any]) -> None:
if key.startswith("per_channel_statistics"):
new_key = ".".join(["decoder", key])
param = state_dict.pop(key)
state_dict[new_key] = param
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_VAE_SPECIAL_KEYS_REMAP = {
"per_channel_statistics.channel": remove_keys_inplace,
"per_channel_statistics.mean-of-stds": remove_keys_inplace,
}
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]) -> tuple[dict[str, Any], dict[str, Any]]:
connector_prefixes = (
"video_embeddings_connector",
"audio_embeddings_connector",
"transformer_1d_blocks",
"text_embedding_projection",
"connectors.",
"video_connector",
"audio_connector",
"text_proj_in",
)
transformer_state_dict, connector_state_dict = {}, {}
for key, value in state_dict.items():
if key.startswith(connector_prefixes):
connector_state_dict[key] = value
else:
transformer_state_dict[key] = value
return transformer_state_dict, connector_state_dict
def get_ltx2_transformer_config(version: str) -> 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
elif version == "2.5":
config = {
"model_id": "Lightricks/LTX-2.5",
"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,
# The only transformer-level deltas from 2.3: the video FFN drops its bias (audio_ff_bias and
# use_prompt_adaln_single keep their True defaults for this checkpoint), and 2.5 carries a
# learned keyframe absolute-position embedding.
"ff_bias": False,
"use_keyframes_abs_pos_embedding": 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(
version: str, gemma_text_config: Any | None = None
) -> 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
elif version == "2.5":
if gemma_text_config is None:
raise ValueError("gemma_text_config is required to derive connector dims for LTX-2.5.")
config = {
"model_id": "Lightricks/LTX-2.5",
"diffusers_config": {
# Derived from the Gemma 4 text config rather than hardcoded, since (unlike Gemma-3-12B) the
# 2.5 text encoder isn't a single fixed checkpoint. Formula matches the reference
# (`encoder_configurator._create_feature_extractor`): hidden_size, and num_hidden_layers + 1
# for the embedding layer.
"caption_channels": gemma_text_config.hidden_size,
"text_proj_in_factor": gemma_text_config.num_hidden_layers + 1,
"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], version: str) -> 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(original_state_dict)
with init_empty_weights():
transformer = LTX2VideoTransformer3DModel.from_config(diffusers_config)
# Handle official code --> diffusers key remapping via the remap dict
for key in list(transformer_state_dict.keys()):
new_key = key[:]
for replace_key, rename_key in rename_dict.items():
new_key = new_key.replace(replace_key, rename_key)
update_state_dict_inplace(transformer_state_dict, key, new_key)
# Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in
# special_keys_remap
for key in list(transformer_state_dict.keys()):
for special_key, handler_fn_inplace in special_keys_remap.items():
if special_key not in key:
continue
handler_fn_inplace(key, transformer_state_dict)
transformer.load_state_dict(transformer_state_dict, strict=True, assign=True)
return transformer
def convert_ltx2_connectors(
original_state_dict: dict[str, Any], version: str, gemma_text_config: Any | None = None
) -> LTX2TextConnectors:
config, rename_dict, special_keys_remap = get_ltx2_connectors_config(version, gemma_text_config=gemma_text_config)
diffusers_config = config["diffusers_config"]
_, connector_state_dict = split_transformer_and_connector_state_dict(original_state_dict)
if len(connector_state_dict) == 0:
raise ValueError("No connector weights found in the provided state dict.")
with init_empty_weights():
connectors = LTX2TextConnectors.from_config(diffusers_config)
for key in list(connector_state_dict.keys()):
new_key = key[:]
for replace_key, rename_key in rename_dict.items():
new_key = new_key.replace(replace_key, rename_key)
update_state_dict_inplace(connector_state_dict, key, new_key)
for key in list(connector_state_dict.keys()):
for special_key, handler_fn_inplace in special_keys_remap.items():
if special_key not in key:
continue
handler_fn_inplace(key, connector_state_dict)
connectors.load_state_dict(connector_state_dict, strict=True, assign=True)
return connectors
def convert_ltx2_duration_head(original_state_dict: dict[str, Any]) -> LTX2DurationHead | None:
"""Builds an `LTX2DurationHead` from a duration-head state dict, or `None` if the checkpoint has none.
The duration head ships from LTX-2.5 onward. Its hyperparameters are absent from checkpoint metadata, so the
dimensions are read back from the weight shapes; `num_pooler_heads` cannot be recovered that way and is fixed at
the value the head was trained with.
The original checkpoint stores the pooler's projections fused, in `torch.nn.MultiheadAttention` layout. They are
split here into the separate q/k/v projections `LTX2DurationAttentionPooler` uses.
"""
state_dict = dict(original_state_dict)
if len(state_dict) == 0:
print("No duration_head weights found in the checkpoint; skipping (expected for pre-2.5 checkpoints).")
return None
in_proj_weight = state_dict.pop("attention_pooler.cross_attn.in_proj_weight")
in_proj_bias = state_dict.pop("attention_pooler.cross_attn.in_proj_bias")
query_weight, key_weight, value_weight = in_proj_weight.chunk(3, dim=0)
query_bias, key_bias, value_bias = in_proj_bias.chunk(3, dim=0)
state_dict["attention_pooler.to_q.weight"] = query_weight
state_dict["attention_pooler.to_k.weight"] = key_weight
state_dict["attention_pooler.to_v.weight"] = value_weight
state_dict["attention_pooler.to_q.bias"] = query_bias
state_dict["attention_pooler.to_k.bias"] = key_bias
state_dict["attention_pooler.to_v.bias"] = value_bias
state_dict["attention_pooler.to_out.weight"] = state_dict.pop("attention_pooler.cross_attn.out_proj.weight")
state_dict["attention_pooler.to_out.bias"] = state_dict.pop("attention_pooler.cross_attn.out_proj.bias")
diffusers_config = {
"video_cross_attention_dim": state_dict["video_input_proj.weight"].shape[1],
"audio_cross_attention_dim": state_dict["audio_input_proj.weight"].shape[1],
"pooler_hidden_dim": state_dict["attention_pooler.to_q.weight"].shape[1],
"num_queries": state_dict["attention_pooler.query_tokens"].shape[0],
"mlp_hidden_dim": state_dict["mlp_hidden.weight"].shape[0],
"num_pooler_heads": 4,
}
with init_empty_weights():
duration_head = LTX2DurationHead.from_config(diffusers_config)
duration_head.load_state_dict(state_dict, strict=True, assign=True)
return duration_head
def get_ltx2_video_vae_config(
version: str, 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
elif version == "2.5":
# Same block structure as 2.3 (32x32x8 compression); confirmed against the checkpoint's
# config["vae"]["encoder_blocks"]/["decoder_blocks"] metadata, which is byte-identical to 2.3's.
config = {
"model_id": "Lightricks/LTX-2.5",
"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], version: str, timestep_conditioning: bool
) -> 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():
vae = AutoencoderKLLTX2Video.from_config(diffusers_config)
# Handle official code --> diffusers key remapping via the remap dict
for key in list(original_state_dict.keys()):
new_key = key[:]
for replace_key, rename_key in rename_dict.items():
new_key = new_key.replace(replace_key, rename_key)
update_state_dict_inplace(original_state_dict, key, new_key)
# Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in
# special_keys_remap
for key in list(original_state_dict.keys()):
for special_key, handler_fn_inplace in special_keys_remap.items():
if special_key not in key:
continue
handler_fn_inplace(key, original_state_dict)
vae.load_state_dict(original_state_dict, strict=True, assign=True)
return vae
def get_ltx2_diffusion_video_vae_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
if version != "2.5":
raise ValueError(
f"The diffusion decoder was introduced in LTX-2.5, which the converter handles under "
f"`--version 2.5`; got version {version!r}."
)
# The encoder half is 2.5's conv VAE encoder, unchanged (its weights are byte identical to 2.3's), so
# those entries must stay in sync with `get_ltx2_video_vae_config("2.5")`.
config = {
"model_id": "Lightricks/LTX-2.5",
"diffusers_config": {
"out_channels": 3,
"latent_channels": 128,
"patch_size": 4,
"decoder_head_dim": 64,
"decoder_stage_channels": (2048, 1024, 512, 512, 256),
"decoder_stage_depths": (4, 6, 4, 2, 8),
"decoder_stage_kernels": ((3, 7, 7), (3, 7, 7), (3, 5, 5), (3, 5, 5)),
"decoder_upsample_strides": ((1, 2, 2), (2, 1, 1), (2, 2, 2), (2, 2, 2)),
"decoder_upsample_channel_reductions": (2, 2, 1, 2),
"decoder_stage5_kernel": (11, 11, 11),
"decoder_t_emb_dim": 384,
"decoder_timestep_scale_multiplier": 1000.0,
"decoder_model_output_type": "x0",
"decoder_num_inference_steps": 1,
"spatial_compression_ratio": 32,
"temporal_compression_ratio": 8,
},
}
return config, LTX_2_3_VIDEO_VAE_RENAME_DICT, LTX_2_0_VAE_SPECIAL_KEYS_REMAP
def convert_ltx2_diffusion_video_vae(original_state_dict: dict[str, Any], version: str) -> dict[str, Any]:
config, rename_dict, special_keys_remap = get_ltx2_diffusion_video_vae_config(version)
diffusers_config = config["diffusers_config"]
with init_empty_weights():
vae = LTX2VideoDiffusionDecoderModel.from_config(diffusers_config)
# The checkpoint is a whole VAE, but this model is decoder-only: encoding stays with
# `AutoencoderKLLTX2Video`. So keep the `decoder.` half and the per-channel statistics (which become
# `latents_mean` / `latents_std`), and drop the encoder outright rather than remapping weights that
# have nowhere to land. `load_state_dict` below is strict, so anything left over would raise.
decoder_state_dict = {
key.removeprefix("decoder."): value for key, value in original_state_dict.items() if key.startswith("decoder.")
}
for key in list(original_state_dict.keys()):
if key.startswith("decoder.") or not key.startswith("per_channel_statistics"):
del original_state_dict[key]
for key in list(original_state_dict.keys()):
new_key = key[:]
for replace_key, rename_key in rename_dict.items():
new_key = new_key.replace(replace_key, rename_key)
update_state_dict_inplace(original_state_dict, key, new_key)
for key in list(original_state_dict.keys()):
for special_key, handler_fn_inplace in special_keys_remap.items():
if special_key not in key:
continue
handler_fn_inplace(key, original_state_dict)
gates = {
key: value
for key, value in decoder_state_dict.items()
if key.endswith((".gate_msa", ".gate_mlp", ".gate_ctx"))
}
for key, value in decoder_state_dict.items():
# Bundled preview heads are not part of the decoder. The sft (dev) checkpoint carries neither these
# nor any gate, so both this and the fold below no-op on it; gated distilled checkpoints need them.
if key.startswith("coarse_") or ".coarse_" in key or key in gates:
continue
new_key = key[:]
for replace_key, rename_key in LTX_2_5_DIFFUSION_DECODER_RENAME_DICT.items():
new_key = new_key.replace(replace_key, rename_key)
gated_leaf = next(
(leaf for leaf in LTX_2_5_DIFFUSION_DECODER_GATE_FOLD_TARGETS if new_key.endswith(leaf)), None
)
if gated_leaf is not None:
# The gate is a sibling of the Linear it gates, so it shares the block prefix.
gate = gates.get(new_key[: -len(gated_leaf)] + LTX_2_5_DIFFUSION_DECODER_GATE_FOLD_TARGETS[gated_leaf])
if gate is not None:
gate = gate.to(torch.float32)
folded = (gate.unsqueeze(1) if value.ndim == 2 else gate) * value.to(torch.float32)
value = folded.to(value.dtype)
# The checkpoint stores one fused `Linear(dim, 3 * dim)`; the model owns three separate projections.
if new_key.endswith((".qkv.weight", ".qkv.bias")):
leaf = "weight" if new_key.endswith(".weight") else "bias"
prefix = new_key[: -len(f"qkv.{leaf}")]
if value.shape[0] % 3 != 0:
raise ValueError(f"Fused QKV param {key!r} leading dim {value.shape[0]} is not divisible by 3.")
chunk = value.shape[0] // 3
original_state_dict[f"decoder.{prefix}to_q.{leaf}"] = value[:chunk].clone()
original_state_dict[f"decoder.{prefix}to_k.{leaf}"] = value[chunk : 2 * chunk].clone()
original_state_dict[f"decoder.{prefix}to_v.{leaf}"] = value[2 * chunk :].clone()
continue
original_state_dict[f"decoder.{new_key}"] = value
vae.load_state_dict(original_state_dict, strict=True, assign=True)
return vae
def get_ltx2_audio_vae_config(version: str) -> 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
elif version == "2.5":
config = {
"model_id": "Lightricks/LTX-2.5",
"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 / 2.3
}
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], version: str) -> dict[str, Any]:
config, rename_dict, special_keys_remap = get_ltx2_audio_vae_config(version)
diffusers_config = config["diffusers_config"]
with init_empty_weights():
vae = AutoencoderKLLTX2Audio.from_config(diffusers_config)
# Handle official code --> diffusers key remapping via the remap dict
for key in list(original_state_dict.keys()):
new_key = key[:]
for replace_key, rename_key in rename_dict.items():
new_key = new_key.replace(replace_key, rename_key)
update_state_dict_inplace(original_state_dict, key, new_key)
# Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in
# special_keys_remap
for key in list(original_state_dict.keys()):
for special_key, handler_fn_inplace in special_keys_remap.items():
if special_key not in key:
continue
handler_fn_inplace(key, original_state_dict)
vae.load_state_dict(original_state_dict, strict=True, assign=True)
return vae
def get_ltx2_vocoder_config(version: str) -> 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
elif version == "2.5":
config = {
"model_id": "Lightricks/LTX-2.5",
"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,
}, # Same config as LTX-2.3
}
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], version: str) -> dict[str, Any]:
config, rename_dict, special_keys_remap = get_ltx2_vocoder_config(version)
diffusers_config = config["diffusers_config"]
if version in ("2.3", "2.5"):
vocoder_cls = LTX2VocoderWithBWE
else:
vocoder_cls = LTX2Vocoder
with init_empty_weights():
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()):
new_key = key[:]
for replace_key, rename_key in rename_dict.items():
new_key = new_key.replace(replace_key, rename_key)
update_state_dict_inplace(original_state_dict, key, new_key)
# Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in
# special_keys_remap
for key in list(original_state_dict.keys()):
for special_key, handler_fn_inplace in special_keys_remap.items():
if special_key not in key:
continue
handler_fn_inplace(key, original_state_dict)
vocoder.load_state_dict(original_state_dict, strict=True, assign=True)
return vocoder
def get_ltx2_spatial_latent_upsampler_config(version: str):
if version == "2.0":
config = {
"in_channels": 128,
"mid_channels": 1024,
"num_blocks_per_stage": 4,
"dims": 3,
"spatial_upsample": True,
"temporal_upsample": False,
"rational_spatial_scale": 2.0,
"use_rational_resampler": True,
}
elif version in ("2.3", "2.5"):
config = {
"in_channels": 128,
"mid_channels": 1024,
"num_blocks_per_stage": 4,
"dims": 3,
"spatial_upsample": True,
"temporal_upsample": False,
"rational_spatial_scale": 2.0,
"use_rational_resampler": False,
}
else:
raise ValueError(f"Unsupported version: {version}")
return config
def get_ltx2_temporal_latent_upsampler_config(version: str):
if version != "2.5":
raise ValueError(f"Unsupported version: {version}")
# The temporal x2 upsampler is narrower than its spatial sibling and pixel-shuffles along time only.
return {
"in_channels": 128,
"mid_channels": 512,
"num_blocks_per_stage": 4,
"dims": 3,
"spatial_upsample": False,
"temporal_upsample": True,
}
def convert_ltx2_latent_upsampler(original_state_dict: dict[str, Any], config: dict[str, Any], dtype: torch.dtype):
with init_empty_weights():
latent_upsampler = LTX2LatentUpsamplerModel(**config)
latent_upsampler.load_state_dict(original_state_dict, strict=True, assign=True)
latent_upsampler.to(dtype)
return latent_upsampler
def load_original_checkpoint(args, filename: str | None) -> dict[str, Any]:
if args.original_state_dict_repo_id is not None:
ckpt_path = hf_hub_download(repo_id=args.original_state_dict_repo_id, filename=filename)
elif args.checkpoint_path is not None:
ckpt_path = args.checkpoint_path
else:
raise ValueError("Please provide either `original_state_dict_repo_id` or a local `checkpoint_path`")
original_state_dict = safetensors.torch.load_file(ckpt_path)
return original_state_dict
def load_hub_or_local_checkpoint(repo_id: str | None = None, filename: str | None = None) -> dict[str, Any]:
if repo_id is None and filename is None:
raise ValueError("Please supply at least one of `repo_id` or `filename`")
if repo_id is not None:
if filename is None:
raise ValueError("If repo_id is specified, filename must also be specified.")
ckpt_path = hf_hub_download(repo_id=repo_id, filename=filename)
else:
ckpt_path = filename
_, ext = os.path.splitext(ckpt_path)
if ext in [".safetensors", ".sft"]:
state_dict = safetensors.torch.load_file(ckpt_path)
else:
state_dict = torch.load(ckpt_path, map_location="cpu")
return state_dict
def get_model_state_dict_from_combined_ckpt(combined_ckpt: dict[str, Any], prefix: str) -> dict[str, Any]:
# Ensure that the key prefix ends with a dot (.)
if not prefix.endswith("."):
prefix = prefix + "."
model_state_dict = {}
for param_name, param in combined_ckpt.items():
if param_name.startswith(prefix):
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_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
def get_args():
parser = argparse.ArgumentParser()
def none_or_str(value: str):
if isinstance(value, str) and value.lower() == "none":
return None
return value
parser.add_argument(
"--original_state_dict_repo_id",
default="Lightricks/LTX-2",
type=none_or_str,
help="HF Hub repo id with LTX 2.0 checkpoint",
)
parser.add_argument(
"--checkpoint_path",
default=None,
type=str,
help="Local checkpoint path for LTX 2.0. Will be used if `original_state_dict_repo_id` is not specified.",
)
parser.add_argument(
"--version",
type=str,
default="2.0",
choices=["test", "2.0", "2.3", "2.5"],
help="Version of the LTX 2.0 model",
)
parser.add_argument(
"--combined_filename",
default="ltx-2-19b-dev.safetensors",
type=none_or_str,
help="Filename for combined checkpoint with all LTX 2.0 models (VAE, DiT, etc.)",
)
parser.add_argument("--vae_prefix", default="vae.", type=str)
parser.add_argument("--audio_vae_prefix", default="audio_vae.", type=str)
parser.add_argument("--dit_prefix", default="model.diffusion_model.", type=str)
parser.add_argument("--vocoder_prefix", default="vocoder.", type=str)
parser.add_argument("--duration_head_prefix", default="duration_head.", type=str)
parser.add_argument("--vae_filename", default=None, type=str, help="VAE filename; overrides combined ckpt if set")
parser.add_argument(
"--audio_vae_filename", default=None, type=str, help="Audio VAE filename; overrides combined ckpt if set"
)
parser.add_argument("--dit_filename", default=None, type=str, help="DiT filename; overrides combined ckpt if set")
parser.add_argument(
"--vocoder_filename", default=None, type=str, help="Vocoder filename; overrides combined ckpt if set"
)
parser.add_argument(
"--text_encoder_model_id",
default="google/gemma-3-12b-it-qat-q4_0-unquantized",
type=none_or_str,
help=(
"HF Hub id for the text encoder model. Default is Gemma 3, used by LTX 2.0/2.3. LTX-2.5 requires a "
"Gemma 4 (`gemma4_unified`) checkpoint here instead -- passing the Gemma 3 default with `--version 2.5` "
"raises an error."
),
)
parser.add_argument(
"--tokenizer_id",
default="google/gemma-3-12b-it-qat-q4_0-unquantized",
type=none_or_str,
help="HF Hub id for the text tokenizer. Should match --text_encoder_model_id's family (Gemma 3 vs Gemma 4).",
)
parser.add_argument(
"--prompt_enhancer_model_id",
default=None,
type=none_or_str,
help=(
"HF Hub id for the prompt-enhancer model (used with --add_processor). For LTX-2.0/2.3, defaults to "
"--text_encoder_model_id (the same Gemma 3 checkpoint serves both roles). LTX-2.5's fine-tuned text "
"encoder is not trained for enhancement, so this must be set explicitly for --version 2.5 -- e.g. to "
"google/gemma-4-E2B-it or google/gemma-4-E4B-it."
),
)
parser.add_argument(
"--temporal_latent_upsampler_filename",
default="ltx-2.5-latent-temporal-upscaler-x2-bf16-1.0.safetensors",
type=none_or_str,
help="Temporal x2 latent upsampler filename (LTX-2.5, used by the DFR pipeline's temporal refine rounds)",
)
parser.add_argument(
"--latent_upsampler_filename",
default="ltx-2-spatial-upscaler-x2-1.0.safetensors",
type=none_or_str,
help="Latent upsampler filename",
)
parser.add_argument(
"--timestep_conditioning", action="store_true", help="Whether to add timestep condition to the video VAE model"
)
parser.add_argument("--vae", action="store_true", help="Whether to convert the video VAE model")
parser.add_argument(
"--diffusion_vae",
action="store_true",
help=(
"Whether to convert the LTX-2.5 diffusion decoder, saved to a `diffusion_decoder` subfolder — the "
"component name `LTX2VideoDiffusionDecodePipeline` and the modular blocks resolve it by — so "
"`from_pretrained` keeps returning the conv decoder in `vae` by default"
),
)
parser.add_argument("--audio_vae", action="store_true", help="Whether to convert the audio VAE model")
parser.add_argument("--dit", action="store_true", help="Whether to convert the DiT model")
parser.add_argument("--connectors", action="store_true", help="Whether to convert the connector model")
parser.add_argument(
"--duration_head",
action="store_true",
help="Whether to convert the duration head (present in LTX-2.5 and later checkpoints only)",
)
parser.add_argument("--vocoder", action="store_true", help="Whether to convert the vocoder model")
parser.add_argument("--text_encoder", action="store_true", help="Whether to conver the text encoder")
parser.add_argument("--latent_upsampler", action="store_true", help="Whether to convert the latent upsampler")
parser.add_argument(
"--temporal_latent_upsampler",
action="store_true",
help="Whether to convert the temporal x2 latent upsampler (LTX-2.5)",
)
parser.add_argument(
"--full_pipeline",
action="store_true",
help="Whether to save the pipeline. This will attempt to convert all models (e.g. vae, dit, etc.)",
)
parser.add_argument(
"--upsample_pipeline",
action="store_true",
help="Whether to save a latent upsampling pipeline",
)
parser.add_argument(
"--add_processor",
action="store_true",
help="Whether to add a text-encoder processor to the pipeline for prompt enhancement.",
)
parser.add_argument("--vae_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"])
parser.add_argument("--audio_vae_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"])
parser.add_argument("--dit_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"])
parser.add_argument("--vocoder_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"])
parser.add_argument("--text_encoder_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"])
parser.add_argument("--output_path", type=str, required=True, help="Path where converted model should be saved")
parser.add_argument(
"--upsample_output_path",
type=str,
default=None,
help="Path where converted upsampling pipeline should be saved",
)
return parser.parse_args()
DTYPE_MAPPING = {
"fp32": torch.float32,
"fp16": torch.float16,
"bf16": torch.bfloat16,
}
VARIANT_MAPPING = {
"fp32": None,
"fp16": "fp16",
"bf16": "bf16",
}
def main(args):
vae_dtype = DTYPE_MAPPING[args.vae_dtype]
audio_vae_dtype = DTYPE_MAPPING[args.audio_vae_dtype]
dit_dtype = DTYPE_MAPPING[args.dit_dtype]
vocoder_dtype = DTYPE_MAPPING[args.vocoder_dtype]
text_encoder_dtype = DTYPE_MAPPING[args.text_encoder_dtype]
combined_ckpt = None
load_combined_models = any(
[
args.vae,
args.diffusion_vae,
args.audio_vae,
args.dit,
args.vocoder,
args.connectors,
args.full_pipeline,
args.upsample_pipeline,
]
)
if args.combined_filename is not None and load_combined_models:
combined_ckpt = load_original_checkpoint(args, filename=args.combined_filename)
# LTX-2.5 only works with a Gemma 4 (`gemma4_unified`) text encoder; --text_encoder_model_id defaults to
# Gemma 3 (for 2.0/2.3), so silently proceeding would pair a 2.5 checkpoint with the wrong text encoder.
gemma_text_config = None
if args.version == "2.5" and (args.text_encoder or args.connectors or args.full_pipeline):
gemma_config = AutoConfig.from_pretrained(args.text_encoder_model_id)
if gemma_config.model_type != "gemma4_unified":
raise ValueError(
f"LTX-2.5 requires a Gemma 4 (`gemma4_unified`) text encoder, but --text_encoder_model_id="
f"{args.text_encoder_model_id!r} has model_type={gemma_config.model_type!r}. Pass "
"--text_encoder_model_id pointing at a Gemma 4 checkpoint (the default is Gemma 3, for 2.0/2.3)."
)
gemma_text_config = gemma_config.text_config
# LTX-2.5's fine-tuned text encoder is never a valid prompt enhancer (unlike LTX-2.0/2.3, where the same Gemma 3
# checkpoint serves both roles) -- require an explicit, separate --prompt_enhancer_model_id instead of silently
# falling back to --text_encoder_model_id.
if (
args.version == "2.5"
and args.add_processor
and (args.text_encoder or args.full_pipeline)
and args.prompt_enhancer_model_id is None
):
raise ValueError(
"LTX-2.5's text encoder is not trained for prompt enhancement, so --prompt_enhancer_model_id must be "
"set explicitly when --add_processor is used with --version 2.5 -- e.g. to google/gemma-4-E2B-it or "
"google/gemma-4-E4B-it."
)
if args.vae or args.full_pipeline or args.upsample_pipeline:
if args.vae_filename is not None:
original_vae_ckpt = load_hub_or_local_checkpoint(filename=args.vae_filename)
elif combined_ckpt is not None:
original_vae_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.vae_prefix)
vae = convert_ltx2_video_vae(
original_vae_ckpt, version=args.version, timestep_conditioning=args.timestep_conditioning
)
if not args.full_pipeline and not args.upsample_pipeline:
vae.to(vae_dtype).save_pretrained(os.path.join(args.output_path, "vae"))
if args.diffusion_vae:
if args.vae_filename is not None:
original_diffusion_vae_ckpt = load_hub_or_local_checkpoint(filename=args.vae_filename)
elif combined_ckpt is not None:
original_diffusion_vae_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.vae_prefix)
diffusion_vae = convert_ltx2_diffusion_video_vae(original_diffusion_vae_ckpt, version=args.version)
# "diffusion_decoder", not "vae_diffusion": pipeline-level `from_pretrained` resolves each component
# from the subfolder named after it, so this folder name must match the `diffusion_decoder` component
# of `LTX2VideoDiffusionDecodePipeline` (and the modular `ComponentSpec`) for those loads to work.
diffusion_vae.to(vae_dtype).save_pretrained(os.path.join(args.output_path, "diffusion_decoder"))
if args.audio_vae or args.full_pipeline:
if args.audio_vae_filename is not None:
original_audio_vae_ckpt = load_hub_or_local_checkpoint(filename=args.audio_vae_filename)
elif combined_ckpt is not None:
original_audio_vae_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.audio_vae_prefix)
audio_vae = convert_ltx2_audio_vae(original_audio_vae_ckpt, version=args.version)
if not args.full_pipeline:
audio_vae.to(audio_vae_dtype).save_pretrained(os.path.join(args.output_path, "audio_vae"))
if args.dit or args.full_pipeline:
if args.dit_filename is not None:
original_dit_ckpt = load_hub_or_local_checkpoint(filename=args.dit_filename)
elif combined_ckpt is not None:
original_dit_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.dit_prefix)
transformer = convert_ltx2_transformer(original_dit_ckpt, version=args.version)
if not args.full_pipeline:
transformer.to(dit_dtype).save_pretrained(os.path.join(args.output_path, "transformer"))
if args.connectors or args.full_pipeline:
if args.dit_filename is not None:
original_connectors_ckpt = load_hub_or_local_checkpoint(filename=args.dit_filename)
elif combined_ckpt is not None:
original_connectors_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.dit_prefix)
connectors = convert_ltx2_connectors(
original_connectors_ckpt, version=args.version, gemma_text_config=gemma_text_config
)
if not args.full_pipeline:
connectors.to(dit_dtype).save_pretrained(os.path.join(args.output_path, "connectors"))
duration_head = None
if args.duration_head or args.full_pipeline:
if combined_ckpt is not None:
original_duration_head_ckpt = get_model_state_dict_from_combined_ckpt(
combined_ckpt, args.duration_head_prefix
)
duration_head = convert_ltx2_duration_head(original_duration_head_ckpt)
if duration_head is not None and not args.full_pipeline:
duration_head.to(dit_dtype).save_pretrained(os.path.join(args.output_path, "duration_head"))
if args.vocoder or args.full_pipeline:
if args.vocoder_filename is not None:
original_vocoder_ckpt = load_hub_or_local_checkpoint(filename=args.vocoder_filename)
elif combined_ckpt is not None:
original_vocoder_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.vocoder_prefix)
vocoder = convert_ltx2_vocoder(original_vocoder_ckpt, version=args.version)
if not args.full_pipeline:
vocoder.to(vocoder_dtype).save_pretrained(os.path.join(args.output_path, "vocoder"))
if args.text_encoder or args.full_pipeline:
text_encoder = AutoModelForImageTextToText.from_pretrained(args.text_encoder_model_id)
if not args.full_pipeline:
text_encoder.to(text_encoder_dtype).save_pretrained(os.path.join(args.output_path, "text_encoder"))
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_id)
if not args.full_pipeline:
tokenizer.save_pretrained(os.path.join(args.output_path, "tokenizer"))
processor = None
prompt_enhancer = None
if args.add_processor:
enhancer_model_id = args.prompt_enhancer_model_id or args.text_encoder_model_id
processor = AutoProcessor.from_pretrained(enhancer_model_id)
if not args.full_pipeline:
processor.save_pretrained(os.path.join(args.output_path, "processor"))
if args.prompt_enhancer_model_id is not None:
# Separate, dedicated enhancer model (required for LTX-2.5); for LTX-2.0/2.3, the same
# `text_encoder` checkpoint already serves as its own enhancer, so nothing extra is saved.
prompt_enhancer = AutoModelForImageTextToText.from_pretrained(enhancer_model_id)
if not args.full_pipeline:
prompt_enhancer.to(text_encoder_dtype).save_pretrained(
os.path.join(args.output_path, "prompt_enhancer")
)
if args.latent_upsampler or args.upsample_pipeline:
original_latent_upsampler_ckpt = load_hub_or_local_checkpoint(
repo_id=args.original_state_dict_repo_id, filename=args.latent_upsampler_filename
)
latent_upsampler_config = get_ltx2_spatial_latent_upsampler_config(args.version)
latent_upsampler = convert_ltx2_latent_upsampler(
original_latent_upsampler_ckpt,
latent_upsampler_config,
dtype=vae_dtype,
)
if not args.full_pipeline and not args.upsample_pipeline:
latent_upsampler.save_pretrained(os.path.join(args.output_path, "latent_upsampler"))
if args.temporal_latent_upsampler:
original_temporal_upsampler_ckpt = load_hub_or_local_checkpoint(
repo_id=args.original_state_dict_repo_id, filename=args.temporal_latent_upsampler_filename
)
temporal_latent_upsampler = convert_ltx2_latent_upsampler(
original_temporal_upsampler_ckpt,
get_ltx2_temporal_latent_upsampler_config(args.version),
dtype=vae_dtype,
)
temporal_latent_upsampler.save_pretrained(os.path.join(args.output_path, "temporal_latent_upsampler"))
if args.full_pipeline:
is_distilled_ckpt = "distilled" in args.combined_filename
if is_distilled_ckpt:
# Disable dynamic shifting and terminal shift so that distilled sigmas are used as-is
scheduler = FlowMatchEulerDiscreteScheduler(
use_dynamic_shifting=False,
base_shift=0.95,
max_shift=2.05,
base_image_seq_len=1024,
max_image_seq_len=4096,
shift_terminal=None,
)
else:
scheduler = FlowMatchEulerDiscreteScheduler(
use_dynamic_shifting=True,
base_shift=0.95,
max_shift=2.05,
base_image_seq_len=1024,
max_image_seq_len=4096,
shift_terminal=0.1,
)
pipe = LTX2Pipeline(
scheduler=scheduler,
vae=vae,
audio_vae=audio_vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
connectors=connectors,
transformer=transformer,
vocoder=vocoder,
processor=processor,
prompt_enhancer=prompt_enhancer,
)
pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")
if args.upsample_pipeline:
pipe = LTX2LatentUpsamplePipeline(vae=vae, latent_upsampler=latent_upsampler)
# As two diffusers pipelines cannot be in the same directory, save the upsampling pipeline to its own directory
if args.upsample_output_path:
upsample_output_path = args.upsample_output_path
else:
upsample_output_path = args.output_path
pipe.save_pretrained(upsample_output_path, safe_serialization=True, max_shard_size="5GB")
if __name__ == "__main__":
args = get_args()
main(args)
View Git Blame Copy Permalink