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()):