diff --git a/src/diffusers/loaders/single_file_model.py b/src/diffusers/loaders/single_file_model.py index 610e377d3ab9..3545fa7598ba 100644 --- a/src/diffusers/loaders/single_file_model.py +++ b/src/diffusers/loaders/single_file_model.py @@ -31,6 +31,7 @@ from ..utils.torch_utils import empty_device_cache from .single_file_utils import ( SingleFileComponentError, + convert_anima_text_conditioner_to_diffusers, convert_animatediff_checkpoint_to_diffusers, convert_auraflow_transformer_checkpoint_to_diffusers, convert_autoencoder_dc_checkpoint_to_diffusers, @@ -193,6 +194,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 e65e0be50f62..26005552e387 100644 --- a/src/diffusers/loaders/single_file_utils.py +++ b/src/diffusers/loaders/single_file_utils.py @@ -3736,6 +3736,22 @@ def swap_scale_shift(weight): return converted_state_dict +def convert_anima_text_conditioner_to_diffusers(checkpoint, **kwargs): + # If none of the keys has the `llm_adapter` prefix, we assume it is a diffusers-format file + # containing ONLY the text conditioner weights and nothing else + has_prefix = any(key.startswith("llm_adapter") for key in checkpoint.keys()) + if not has_prefix: + return checkpoint + + converted_state_dict = {} + for k in list(checkpoint.keys()): + if "llm_adapter." in k: + converted_state_dict[k.split("llm_adapter.", 1)[1]] = 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..65ab281d9317 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 FromOriginalModelMixin, PeftAdapterMixin 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.