From 92e463b8d0235948f44d2a28658083f0779d8c86 Mon Sep 17 00:00:00 2001 From: Functionhx <2994114386@qq.com> Date: Fri, 10 Jul 2026 03:12:44 +0800 Subject: [PATCH] Fix LTX2 from_single_file: missing 2.3 key mappings and shared checkpoint dict draining Two independent bugs in the LTX2 single-file converters: Bug 1: Missing LTX 2.3 key mappings. - convert_ltx2_transformer_to_diffusers did not map the 2.3 prompt_adaln_single and audio_prompt_adaln_single modulation keys, leaving prompt_adaln/audio_prompt_adaln modules as meta tensors. - convert_ltx2_vae_to_diffusers did not map the 2.3 decoder up_blocks.7/up_blocks.8 keys. Bug 2: Each LTX2 converter consumed the entire checkpoint dict via checkpoint.pop(), so only the first component got weights and later components received an empty dict. Fixed by filtering to each converter's own key prefix (model.diffusion_model./vae./audio_vae.), falling back to all keys when the prefix is absent for single-component extraction. Fixes #14145. Test Plan: Verified key parity for both LTX 2.0 and 2.3 checkpoints: - LTX 2.3: transformer 4186/4186, vae 170/170, audio_vae 102/102 - LTX 2.0: transformer 3510/3510, vae 184/184, audio_vae 102/102 No missing or unexpected keys in any component. Signed-off-by: Functionhx <2994114386@qq.com> Signed-off-by: Yuchen Fan --- src/diffusers/loaders/single_file_utils.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/src/diffusers/loaders/single_file_utils.py b/src/diffusers/loaders/single_file_utils.py index 296f32f891f0..d301a0a6ca02 100644 --- a/src/diffusers/loaders/single_file_utils.py +++ b/src/diffusers/loaders/single_file_utils.py @@ -4048,6 +4048,18 @@ def convert_ltx2_transformer_adaln_single(key: str, state_dict) -> None: if ".weight" not in key and ".bias" not in key: return + if key.startswith("prompt_adaln_single."): + new_key = key.replace("prompt_adaln_single.", "prompt_adaln.") + param = state_dict.pop(key) + state_dict[new_key] = param + return + + if key.startswith("audio_prompt_adaln_single."): + new_key = key.replace("audio_prompt_adaln_single.", "audio_prompt_adaln.") + param = state_dict.pop(key) + state_dict[new_key] = param + return + if key.startswith("adaln_single."): new_key = key.replace("adaln_single.", "time_embed.") param = state_dict.pop(key) @@ -4066,7 +4078,7 @@ def convert_ltx2_transformer_adaln_single(key: str, state_dict) -> None: "adaln_single": convert_ltx2_transformer_adaln_single, } - converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())} + converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys()) if key.startswith("model.diffusion_model.")} # Handle official code --> diffusers key remapping via the remap dict for key in list(converted_state_dict.keys()): @@ -4109,6 +4121,8 @@ def convert_ltx2_vae_to_diffusers(checkpoint, **kwargs): "up_blocks.4": "up_blocks.1", "up_blocks.5": "up_blocks.2.upsamplers.0", "up_blocks.6": "up_blocks.2", + "up_blocks.7": "up_blocks.3.upsamplers.0", + "up_blocks.8": "up_blocks.3", # Common # For all 3D ResNets "res_blocks": "resnets", @@ -4127,7 +4141,7 @@ def remove_keys_inplace(key: str, state_dict) -> None: "per_channel_statistics.mean-of-stds": remove_keys_inplace, } - converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())} + converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys()) if key.startswith("vae.")} # Handle official code --> diffusers key remapping via the remap dict for key in list(converted_state_dict.keys()): @@ -4159,7 +4173,7 @@ def convert_ltx2_audio_vae_to_diffusers(checkpoint, **kwargs): def update_state_dict_inplace(state_dict, old_key: str, new_key: str) -> None: state_dict[new_key] = state_dict.pop(old_key) - converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())} + converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys()) if key.startswith("audio_vae.")} # Handle official code --> diffusers key remapping via the remap dict for key in list(converted_state_dict.keys()):