diff --git a/src/diffusers/loaders/single_file_model.py b/src/diffusers/loaders/single_file_model.py index 1556227673f4..1a86bc1338ff 100644 --- a/src/diffusers/loaders/single_file_model.py +++ b/src/diffusers/loaders/single_file_model.py @@ -36,6 +36,7 @@ convert_autoencoder_dc_checkpoint_to_diffusers, convert_chroma_transformer_checkpoint_to_diffusers, convert_controlnet_checkpoint, + convert_anima_text_conditioner_to_diffusers, convert_cosmos_transformer_checkpoint_to_diffusers, convert_ernie_image_transformer_checkpoint_to_diffusers, convert_flux2_transformer_checkpoint_to_diffusers, @@ -192,6 +193,10 @@ "checkpoint_mapping_fn": convert_cosmos_transformer_checkpoint_to_diffusers, "default_subfolder": "transformer", }, + "AnimaTextConditioner": { + "checkpoint_mapping_fn": convert_anima_text_conditioner_to_diffusers, + "default_subfolder": "text_conditioner", + }, "QwenImageTransformer2DModel": { "checkpoint_mapping_fn": lambda checkpoint, **kwargs: checkpoint, "default_subfolder": "transformer", diff --git a/src/diffusers/loaders/single_file_utils.py b/src/diffusers/loaders/single_file_utils.py index 3e160be02b4b..6d8ac4f13ad7 100644 --- a/src/diffusers/loaders/single_file_utils.py +++ b/src/diffusers/loaders/single_file_utils.py @@ -3721,6 +3721,14 @@ def swap_scale_shift(weight): return converted_state_dict +def convert_anima_text_conditioner_to_diffusers(checkpoint, **kwargs): + converted_state_dict = {} + for k in list(checkpoint.keys()): + if "llm_adapter" in k: + converted_state_dict[k.removeprefix("net.").removeprefix("model.diffusion_model.").removeprefix("llm_adapter.")] = checkpoint.pop(k) + else: + _ = checkpoint.pop(k) + return converted_state_dict def convert_cosmos_transformer_checkpoint_to_diffusers(checkpoint, **kwargs): converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())} diff --git a/src/diffusers/models/condition_embedders/condition_embedder_anima.py b/src/diffusers/models/condition_embedders/condition_embedder_anima.py index 40fda447ec68..71de1b8fe05f 100644 --- a/src/diffusers/models/condition_embedders/condition_embedder_anima.py +++ b/src/diffusers/models/condition_embedders/condition_embedder_anima.py @@ -17,7 +17,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import PeftAdapterMixin +from ...loaders import PeftAdapterMixin, FromOriginalModelMixin from ..attention import AttentionModuleMixin from ..attention_dispatch import dispatch_attention_fn from ..modeling_utils import ModelMixin @@ -226,7 +226,7 @@ def forward( return hidden_states -class AnimaTextConditioner(ModelMixin, ConfigMixin, PeftAdapterMixin): +class AnimaTextConditioner(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin): r""" Text conditioner used by Anima to map Qwen3 hidden states and T5 token ids to Cosmos text embeddings.