Add ace step lora support

src/diffusers/loaders/__init__.py

@@ -62,6 +62,7 @@ def text_encoder_attn_modules(text_encoder):
     if is_transformers_available():
         _import_structure["single_file"] = ["FromSingleFileMixin"]
         _import_structure["lora_pipeline"] = [
+            "AceStepLoraLoaderMixin",
             "AmusedLoraLoaderMixin",
             "AnimaLoraLoaderMixin",
             "StableDiffusionLoraLoaderMixin",
@@ -118,6 +119,7 @@ def text_encoder_attn_modules(text_encoder):
                 SD3IPAdapterMixin,
             )
             from .lora_pipeline import (
+                AceStepLoraLoaderMixin,
                 AmusedLoraLoaderMixin,
                 AnimaLoraLoaderMixin,
                 AuraFlowLoraLoaderMixin,

src/diffusers/loaders/lora_conversion_utils.py

@@ -3094,3 +3094,33 @@ def convert_module(module):
         raise ValueError(f"`state_dict` should be empty at this point but has {state_dict.keys()=}")
 
     return converted_state_dict
+
+
+def _convert_non_diffusers_ace_step_lora_to_diffusers(state_dict):
+    """Convert an ACE-Step-1.5 (PEFT format) LoRA state dict to diffusers key names.
+
+    The original ACE-Step repo targets ``q_proj``, ``k_proj``, ``v_proj``, ``o_proj``
+    on the DiT decoder while diffusers renames them to ``to_q``, ``to_k``, ``to_v``,
+    ``to_out.0``.  Keys arrive as
+    ``base_model.model.layers.{i}.{self_attn|cross_attn}.{proj}.lora_{A|B}.weight``
+    and are mapped to
+    ``transformer.layers.{i}.{self_attn|cross_attn}.{proj_diffusers}.lora_{A|B}.weight``.
+    """
+    _PROJ_RENAMES = {
+        ".q_proj.": ".to_q.",
+        ".k_proj.": ".to_k.",
+        ".v_proj.": ".to_v.",
+        ".o_proj.": ".to_out.0.",
+    }
+
+    converted_state_dict = {}
+    for key in list(state_dict.keys()):
+        new_key = key
+        if new_key.startswith("base_model.model."):
+            new_key = new_key[len("base_model.model.") :]
+        for old, new in _PROJ_RENAMES.items():
+            new_key = new_key.replace(old, new)
+        new_key = f"transformer.{new_key}"
+        converted_state_dict[new_key] = state_dict.pop(key)
+
+    return converted_state_dict

src/diffusers/loaders/lora_pipeline.py

@@ -46,6 +46,7 @@
     _convert_kohya_flux2_lora_to_diffusers,
     _convert_kohya_flux_lora_to_diffusers,
     _convert_musubi_wan_lora_to_diffusers,
+    _convert_non_diffusers_ace_step_lora_to_diffusers,
     _convert_non_diffusers_anima_lora_to_diffusers,
     _convert_non_diffusers_flux2_lora_to_diffusers,
     _convert_non_diffusers_hidream_lora_to_diffusers,
@@ -6844,6 +6845,208 @@ def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
         super().unfuse_lora(components=components, **kwargs)
 
 
+class AceStepLoraLoaderMixin(LoraBaseMixin):
+    r"""
+    Load LoRA layers into [`AceStepTransformer1DModel`]. Specific to [`AceStepPipeline`].
+    """
+
+    _lora_loadable_modules = ["transformer"]
+    transformer_name = TRANSFORMER_NAME
+
+    @classmethod
+    @validate_hf_hub_args
+    def lora_state_dict(
+        cls,
+        pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+        **kwargs,
+    ):
+        r"""
+        See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details.
+        """
+        cache_dir = kwargs.pop("cache_dir", None)
+        force_download = kwargs.pop("force_download", False)
+        proxies = kwargs.pop("proxies", None)
+        local_files_only = kwargs.pop("local_files_only", None)
+        token = kwargs.pop("token", None)
+        revision = kwargs.pop("revision", None)
+        subfolder = kwargs.pop("subfolder", None)
+        weight_name = kwargs.pop("weight_name", None)
+        use_safetensors = kwargs.pop("use_safetensors", None)
+        return_lora_metadata = kwargs.pop("return_lora_metadata", False)
+
+        allow_pickle = False
+        if use_safetensors is None:
+            use_safetensors = True
+            allow_pickle = True
+
+        user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"}
+
+        state_dict, metadata = _fetch_state_dict(
+            pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
+            weight_name=weight_name,
+            use_safetensors=use_safetensors,
+            local_files_only=local_files_only,
+            cache_dir=cache_dir,
+            force_download=force_download,
+            proxies=proxies,
+            token=token,
+            revision=revision,
+            subfolder=subfolder,
+            user_agent=user_agent,
+            allow_pickle=allow_pickle,
+        )
+
+        is_dora_scale_present = any("dora_scale" in k for k in state_dict)
+        if is_dora_scale_present:
+            warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new."
+            logger.warning(warn_msg)
+            state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
+
+        # Detect original ACE-Step-1.5 PEFT format (q_proj/k_proj naming).
+        is_original_ace_step = any("q_proj" in k or "k_proj" in k for k in state_dict)
+        if is_original_ace_step:
+            state_dict = _convert_non_diffusers_ace_step_lora_to_diffusers(state_dict)
+
+        out = (state_dict, metadata) if return_lora_metadata else state_dict
+        return out
+
+    # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
+    def load_lora_weights(
+        self,
+        pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+        adapter_name: str | None = None,
+        hotswap: bool = False,
+        **kwargs,
+    ):
+        """
+        See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details.
+        """
+        if not USE_PEFT_BACKEND:
+            raise ValueError("PEFT backend is required for this method.")
+
+        low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
+        if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
+            raise ValueError(
+                "`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
+            )
+
+        # if a dict is passed, copy it instead of modifying it inplace
+        if isinstance(pretrained_model_name_or_path_or_dict, dict):
+            pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
+
+        # First, ensure that the checkpoint is a compatible one and can be successfully loaded.
+        kwargs["return_lora_metadata"] = True
+        state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
+
+        is_correct_format = all("lora" in key for key in state_dict.keys())
+        if not is_correct_format:
+            raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.")
+
+        self.load_lora_into_transformer(
+            state_dict,
+            transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer,
+            adapter_name=adapter_name,
+            metadata=metadata,
+            _pipeline=self,
+            low_cpu_mem_usage=low_cpu_mem_usage,
+            hotswap=hotswap,
+        )
+
+    @classmethod
+    # Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.load_lora_into_transformer with SD3Transformer2DModel->AceStepTransformer1DModel
+    def load_lora_into_transformer(
+        cls,
+        state_dict,
+        transformer,
+        adapter_name=None,
+        _pipeline=None,
+        low_cpu_mem_usage=False,
+        hotswap: bool = False,
+        metadata=None,
+    ):
+        """
+        See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details.
+        """
+        if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
+            raise ValueError(
+                "`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
+            )
+
+        # Load the layers corresponding to transformer.
+        logger.info(f"Loading {cls.transformer_name}.")
+        transformer.load_lora_adapter(
+            state_dict,
+            network_alphas=None,
+            adapter_name=adapter_name,
+            metadata=metadata,
+            _pipeline=_pipeline,
+            low_cpu_mem_usage=low_cpu_mem_usage,
+            hotswap=hotswap,
+        )
+
+    @classmethod
+    # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
+    def save_lora_weights(
+        cls,
+        save_directory: str | os.PathLike,
+        transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None,
+        is_main_process: bool = True,
+        weight_name: str = None,
+        save_function: Callable = None,
+        safe_serialization: bool = True,
+        transformer_lora_adapter_metadata: dict | None = None,
+    ):
+        r"""
+        See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information.
+        """
+        lora_layers = {}
+        lora_metadata = {}
+
+        if transformer_lora_layers:
+            lora_layers[cls.transformer_name] = transformer_lora_layers
+            lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata
+
+        if not lora_layers:
+            raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.")
+
+        cls._save_lora_weights(
+            save_directory=save_directory,
+            lora_layers=lora_layers,
+            lora_metadata=lora_metadata,
+            is_main_process=is_main_process,
+            weight_name=weight_name,
+            save_function=save_function,
+            safe_serialization=safe_serialization,
+        )
+
+    # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
+    def fuse_lora(
+        self,
+        components: list[str] = ["transformer"],
+        lora_scale: float = 1.0,
+        safe_fusing: bool = False,
+        adapter_names: list[str] | None = None,
+        **kwargs,
+    ):
+        r"""
+        See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details.
+        """
+        super().fuse_lora(
+            components=components,
+            lora_scale=lora_scale,
+            safe_fusing=safe_fusing,
+            adapter_names=adapter_names,
+            **kwargs,
+        )
+
+    # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
+    def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs):
+        r"""
+        See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details.
+        """
+        super().unfuse_lora(components=components, **kwargs)
+
+
 class LoraLoaderMixin(StableDiffusionLoraLoaderMixin):
     def __init__(self, *args, **kwargs):
         deprecation_message = "LoraLoaderMixin is deprecated and this will be removed in a future version. Please use `StableDiffusionLoraLoaderMixin`, instead."

src/diffusers/models/transformers/ace_step_transformer.py

@@ -21,7 +21,8 @@
 import torch.nn.functional as F
 
 from ...configuration_utils import ConfigMixin, register_to_config
-from ...utils import logging
+from ...loaders import PeftAdapterMixin
+from ...utils import apply_lora_scale, logging
 from ..attention import AttentionMixin, AttentionModuleMixin
 from ..attention_dispatch import (
     AttentionBackendName,
@@ -428,7 +429,7 @@ def forward(
 # --------------------------------------------------------------------------- #
 
 
-class AceStepTransformer1DModel(ModelMixin, ConfigMixin, AttentionMixin, CacheMixin):
+class AceStepTransformer1DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, AttentionMixin, CacheMixin):
     """Diffusion Transformer for ACE-Step 1.5 music generation.
 
     Generates audio latents conditioned on text, lyrics, and timbre. Uses 1D patch embedding (`Conv1d` with stride
@@ -528,13 +529,15 @@ def __init__(
 
         self.gradient_checkpointing = False
 
+    @apply_lora_scale("attention_kwargs")
     def forward(
         self,
         hidden_states: torch.Tensor,
         timestep: torch.Tensor,
         timestep_r: torch.Tensor,
         encoder_hidden_states: torch.Tensor,
         context_latents: torch.Tensor,
+        attention_kwargs: Optional[dict] = None,
         return_dict: bool = True,
     ) -> Union[torch.Tensor, Transformer2DModelOutput]:
         """The [`AceStepTransformer1DModel`] forward method.
@@ -551,6 +554,9 @@ def forward(
             context_latents (`torch.Tensor` of shape `(batch_size, seq_len, context_dim)`):
                 Context latents (source latents concatenated with chunk masks) — fed to the patchify conv alongside
                 `hidden_states`.
+            attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary passed along to the `AttentionProcessor`. Used to pass the LoRA scale via
+                `{"scale": float}`.
             return_dict (`bool`, defaults to `True`):
                 Whether to return a `Transformer2DModelOutput` or a plain tuple.
 

src/diffusers/pipelines/ace_step/pipeline_ace_step.py

@@ -20,6 +20,7 @@
 from transformers import PreTrainedModel, PreTrainedTokenizerFast
 
 from ...guiders.adaptive_projected_guidance import MomentumBuffer, normalized_guidance
+from ...loaders import AceStepLoraLoaderMixin
 from ...models import AutoencoderOobleck
 from ...models.transformers.ace_step_transformer import AceStepTransformer1DModel
 from ...schedulers import FlowMatchEulerDiscreteScheduler
@@ -129,7 +130,7 @@ def _normalize_audio_codes(audio_codes: Union[str, List[str]], batch_size: int)
 """
 
 
-class AceStepPipeline(DiffusionPipeline):
+class AceStepPipeline(DiffusionPipeline, AceStepLoraLoaderMixin):
     r"""
     Pipeline for text-to-music generation using ACE-Step 1.5.
 
@@ -224,6 +225,10 @@ def guidance_scale(self) -> float:
     def num_timesteps(self) -> int:
         return self._num_timesteps
 
+    @property
+    def attention_kwargs(self):
+        return self._attention_kwargs
+
     def check_inputs(
         self,
         prompt: Union[str, List[str]],
@@ -825,6 +830,7 @@ def __call__(
         cfg_interval_start: float = 0.0,
         cfg_interval_end: float = 1.0,
         timesteps: Optional[List[float]] = None,
+        attention_kwargs: Optional[dict] = None,
     ):
         r"""
         The call function to the pipeline for music generation.
@@ -983,6 +989,7 @@ def __call__(
         # step-end callback can read them without the full arg bundle.
         self._guidance_scale = guidance_scale
         self._num_timesteps = num_inference_steps
+        self._attention_kwargs = attention_kwargs
         self._interrupt = False
 
         # Auto-generate instruction based on task_type if not provided
@@ -1176,6 +1183,7 @@ def __call__(
                         timestep_r=torch.cat([t_curr_tensor, t_curr_tensor], dim=0),
                         encoder_hidden_states=torch.cat([encoder_hidden_states, null_encoder_hidden_states], dim=0),
                         context_latents=torch.cat([context_latents, context_latents], dim=0),
+                        attention_kwargs=self.attention_kwargs,
                         return_dict=False,
                     )
                     vt_cond, vt_uncond = model_output[0].chunk(2, dim=0)
@@ -1199,6 +1207,7 @@ def __call__(
                         timestep_r=t_curr_tensor,
                         encoder_hidden_states=encoder_hidden_states,
                         context_latents=context_latents,
+                        attention_kwargs=self.attention_kwargs,
                         return_dict=False,
                     )
                     vt = model_output[0]
@@ -1211,6 +1220,7 @@ def __call__(
                         timestep_r=t_curr_tensor,
                         encoder_hidden_states=non_cover_encoder_hidden_states,
                         context_latents=context_latents,
+                        attention_kwargs=self.attention_kwargs,
                         return_dict=False,
                     )
                     vt_nc = nc_output[0]