From ab200465cbe7576dd7a4eafa785b72d81735f178 Mon Sep 17 00:00:00 2001 From: zcy Date: Wed, 15 Jul 2026 17:30:23 +0800 Subject: [PATCH 1/9] feat: add LoRA support to AceStepPipeline Add standard diffusers LoRA API (load_lora_weights, set_adapters, fuse_lora, save_lora_weights) for AceStepPipeline, targeting the transformer (AceStepTransformer1DModel) only. Includes automatic conversion of ACE-Step-1.5 original PEFT format LoRA weights (q_proj/k_proj/v_proj/o_proj naming) to diffusers format (to_q/to_k/to_v/to_out.0 naming) during loading. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/diffusers/loaders/__init__.py | 2 + .../loaders/lora_conversion_utils.py | 30 +++ src/diffusers/loaders/lora_pipeline.py | 203 +++++++++++++++++ .../transformers/ace_step_transformer.py | 3 +- .../pipelines/ace_step/pipeline_ace_step.py | 3 +- tests/lora/test_lora_layers_ace_step.py | 214 ++++++++++++++++++ 6 files changed, 453 insertions(+), 2 deletions(-) create mode 100644 tests/lora/test_lora_layers_ace_step.py diff --git a/src/diffusers/loaders/__init__.py b/src/diffusers/loaders/__init__.py index 1b0661d4c251..1c6693bd0c08 100644 --- a/src/diffusers/loaders/__init__.py +++ b/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, diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index e1bc530eb29d..d284cb5d23c9 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/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 diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py index 81ebd2f81102..ee1f784b1983 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/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." diff --git a/src/diffusers/models/transformers/ace_step_transformer.py b/src/diffusers/models/transformers/ace_step_transformer.py index 3430d347606a..83e66e61b1de 100644 --- a/src/diffusers/models/transformers/ace_step_transformer.py +++ b/src/diffusers/models/transformers/ace_step_transformer.py @@ -21,6 +21,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import PeftAdapterMixin from ...utils import logging from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import ( @@ -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 diff --git a/src/diffusers/pipelines/ace_step/pipeline_ace_step.py b/src/diffusers/pipelines/ace_step/pipeline_ace_step.py index 26c14d8bfac7..95800309ade2 100644 --- a/src/diffusers/pipelines/ace_step/pipeline_ace_step.py +++ b/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. diff --git a/tests/lora/test_lora_layers_ace_step.py b/tests/lora/test_lora_layers_ace_step.py new file mode 100644 index 000000000000..32c743a8b15a --- /dev/null +++ b/tests/lora/test_lora_layers_ace_step.py @@ -0,0 +1,214 @@ +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import sys +import unittest + +import torch +from transformers import AutoTokenizer, Qwen3Config, Qwen3Model + +from diffusers import AutoencoderOobleck, FlowMatchEulerDiscreteScheduler +from diffusers.models.transformers.ace_step_transformer import AceStepTransformer1DModel +from diffusers.pipelines.ace_step import AceStepConditionEncoder, AceStepPipeline +from diffusers.utils.import_utils import is_peft_available + +from ..testing_utils import ( + require_peft_backend, + skip_mps, +) + + +if is_peft_available(): + from peft import LoraConfig + +sys.path.append(".") + +from .utils import PeftLoraLoaderMixinTests # noqa: E402 + + +@require_peft_backend +@skip_mps +class AceStepLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): + pipeline_class = AceStepPipeline + scheduler_cls = FlowMatchEulerDiscreteScheduler + scheduler_kwargs = {"num_train_timesteps": 1, "shift": 1.0} + + transformer_cls = AceStepTransformer1DModel + transformer_kwargs = { + "hidden_size": 32, + "intermediate_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 8, + "in_channels": 24, + "audio_acoustic_hidden_dim": 8, + "patch_size": 2, + "rope_theta": 10000.0, + "sliding_window": 16, + } + + vae_cls = AutoencoderOobleck + vae_kwargs = { + "encoder_hidden_size": 6, + "downsampling_ratios": [1, 2], + "decoder_channels": 3, + "decoder_input_channels": 8, + "audio_channels": 2, + "channel_multiples": [2, 4], + "sampling_rate": 4, + } + + tokenizer_cls, tokenizer_id = AutoTokenizer, "/nas/zcy/github_issue/models/Qwen3-Embedding-0.6B-tokenizer" + text_encoder_cls, text_encoder_id = None, None + + text_encoder_target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"] + supports_text_encoder_loras = False + + @property + def output_shape(self): + return (1, 2, 1) + + def get_dummy_components(self, scheduler_cls=None, use_dora=False, lora_alpha=None): + scheduler_cls = scheduler_cls or self.scheduler_cls + rank = 4 + lora_alpha = rank if lora_alpha is None else lora_alpha + + torch.manual_seed(0) + transformer = self.transformer_cls(**self.transformer_kwargs) + + scheduler = scheduler_cls(**self.scheduler_kwargs) + + torch.manual_seed(0) + vae = self.vae_cls(**self.vae_kwargs) + + torch.manual_seed(0) + qwen3_config = Qwen3Config( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + vocab_size=151936, + max_position_embeddings=256, + ) + text_encoder = Qwen3Model(qwen3_config) + tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_id) + + torch.manual_seed(0) + condition_encoder = AceStepConditionEncoder( + hidden_size=32, + intermediate_size=64, + text_hidden_dim=32, + timbre_hidden_dim=8, + num_lyric_encoder_hidden_layers=2, + num_timbre_encoder_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + rope_theta=10000.0, + sliding_window=16, + ) + + text_lora_config = LoraConfig( + r=rank, + lora_alpha=lora_alpha, + target_modules=self.text_encoder_target_modules, + init_lora_weights=False, + use_dora=use_dora, + ) + + denoiser_lora_config = LoraConfig( + r=rank, + lora_alpha=lora_alpha, + target_modules=self.denoiser_target_modules, + init_lora_weights=False, + use_dora=use_dora, + ) + + pipeline_components = { + "scheduler": scheduler, + "vae": vae, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "transformer": transformer, + "condition_encoder": condition_encoder, + "audio_tokenizer": None, + "audio_token_detokenizer": None, + } + + return pipeline_components, text_lora_config, denoiser_lora_config + + def get_dummy_inputs(self, with_generator=True): + generator = torch.manual_seed(0) + noise = torch.randn(1, 4, 8) + input_ids = torch.randint(1, 10, size=(1, 10), generator=generator) + + pipeline_inputs = { + "prompt": "A beautiful piano piece", + "lyrics": "[verse]\nSoft notes", + "audio_duration": 0.4, + "num_inference_steps": 2, + "max_text_length": 32, + "output_type": "np", + } + if with_generator: + pipeline_inputs["generator"] = generator + + return noise, input_ids, pipeline_inputs + + @unittest.skip("Not supported in AceStep.") + def test_simple_inference_with_text_denoiser_block_scale(self): + pass + + @unittest.skip("Not supported in AceStep.") + def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): + pass + + @unittest.skip("Not supported in AceStep.") + def test_modify_padding_mode(self): + pass + + @unittest.skip("Not supported in AceStep.") + def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): + pass + + @unittest.skip("Tiny AceStep GQA model produces numerically close outputs for different LoRA ranks.") + def test_correct_lora_configs_with_different_ranks(self): + pass + + @unittest.skip("AceStepPipeline does not accept attention_kwargs for LoRA scale.") + def test_simple_inference_with_text_denoiser_lora_and_scale(self): + pass + + @unittest.skip("AceStepPipeline does not accept attention_kwargs for LoRA scale.") + def test_lora_scale_kwargs_match_fusion(self): + pass + + @unittest.skip("AceStepPipeline does not accept attention_kwargs for LoRA scale.") + def test_set_adapters_match_attention_kwargs(self): + pass + + @unittest.skip("AceStepPipeline does not accept attention_kwargs for LoRA scale.") + def test_simple_inference_with_text_lora_and_scale(self): + pass + + @unittest.skip("AceStep attention layers have no bias; lora_bias is not applicable.") + def test_lora_B_bias(self): + pass + + @unittest.skip("AceStep uses 'layers' not 'transformer_blocks'/'blocks'; base test hardcodes block names.") + def test_lora_fuse_nan(self): + pass From de90a1e7aad75180fc534bc7110e96882ae648e4 Mon Sep 17 00:00:00 2001 From: zcy Date: Wed, 15 Jul 2026 19:25:42 +0800 Subject: [PATCH 2/9] fix: use HuggingFace Hub ID for tokenizer in AceStep LoRA test and apply ruff formatting - Replace hardcoded local path with "Qwen/Qwen3-Embedding-0.6B" so tests run on CI - Apply ruff formatting fix (slice spacing) Co-Authored-By: Claude Opus 4.6 (1M context) --- src/diffusers/loaders/lora_conversion_utils.py | 2 +- tests/lora/test_lora_layers_ace_step.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index d284cb5d23c9..2c22a06f991b 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -3117,7 +3117,7 @@ def _convert_non_diffusers_ace_step_lora_to_diffusers(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."):] + 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}" diff --git a/tests/lora/test_lora_layers_ace_step.py b/tests/lora/test_lora_layers_ace_step.py index 32c743a8b15a..d97454894bac 100644 --- a/tests/lora/test_lora_layers_ace_step.py +++ b/tests/lora/test_lora_layers_ace_step.py @@ -70,7 +70,7 @@ class AceStepLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): "sampling_rate": 4, } - tokenizer_cls, tokenizer_id = AutoTokenizer, "/nas/zcy/github_issue/models/Qwen3-Embedding-0.6B-tokenizer" + tokenizer_cls, tokenizer_id = AutoTokenizer, "Qwen/Qwen3-Embedding-0.6B" text_encoder_cls, text_encoder_id = None, None text_encoder_target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"] From 37d53170a3671b93cddc0e849889ffb9de49ebb2 Mon Sep 17 00:00:00 2001 From: zcy Date: Thu, 16 Jul 2026 15:38:35 +0800 Subject: [PATCH 3/9] fix: override test_lora_fuse_nan instead of skipping it The base test hardcodes block attribute names (transformer_blocks, blocks, etc.) but AceStep uses 'layers'. Override with AceStep-specific block access path instead of skipping. Co-Authored-By: Claude Opus 4.6 (1M context) --- tests/lora/test_lora_layers_ace_step.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/tests/lora/test_lora_layers_ace_step.py b/tests/lora/test_lora_layers_ace_step.py index d97454894bac..edb8ccf36608 100644 --- a/tests/lora/test_lora_layers_ace_step.py +++ b/tests/lora/test_lora_layers_ace_step.py @@ -209,6 +209,23 @@ def test_simple_inference_with_text_lora_and_scale(self): def test_lora_B_bias(self): pass - @unittest.skip("AceStep uses 'layers' not 'transformer_blocks'/'blocks'; base test hardcodes block names.") def test_lora_fuse_nan(self): - pass + import numpy as np + + components, _, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to("cpu") + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) + + pipe.transformer.add_adapter(denoiser_lora_config, "adapter-1") + + with torch.no_grad(): + pipe.transformer.layers[0].self_attn.to_q.lora_A["adapter-1"].weight += float("inf") + + with self.assertRaises(ValueError): + pipe.fuse_lora(safe_fusing=True) + + pipe.fuse_lora(safe_fusing=False) + out = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(np.isnan(out).all()) From 99c3d74ee363151ed1ab702d2d7f242bbc0370d1 Mon Sep 17 00:00:00 2001 From: zcy Date: Thu, 16 Jul 2026 15:46:34 +0800 Subject: [PATCH 4/9] feat: add joint_attention_kwargs to AceStepPipeline for per-step LoRA scale - Add @apply_lora_scale("joint_attention_kwargs") decorator and joint_attention_kwargs param to AceStepTransformer1DModel.forward - Add joint_attention_kwargs param to AceStepPipeline.__call__ and wire it through to all three transformer calls - Un-skip 4 LoRA scale tests that now pass (39 passed, 13 skipped) Co-Authored-By: Claude Opus 4.6 (1M context) --- .../models/transformers/ace_step_transformer.py | 4 +++- .../pipelines/ace_step/pipeline_ace_step.py | 9 +++++++++ tests/lora/test_lora_layers_ace_step.py | 16 ---------------- 3 files changed, 12 insertions(+), 17 deletions(-) diff --git a/src/diffusers/models/transformers/ace_step_transformer.py b/src/diffusers/models/transformers/ace_step_transformer.py index 83e66e61b1de..6ba1f17738d6 100644 --- a/src/diffusers/models/transformers/ace_step_transformer.py +++ b/src/diffusers/models/transformers/ace_step_transformer.py @@ -22,7 +22,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin -from ...utils import logging +from ...utils import apply_lora_scale, logging from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import ( AttentionBackendName, @@ -529,6 +529,7 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("joint_attention_kwargs") def forward( self, hidden_states: torch.Tensor, @@ -536,6 +537,7 @@ def forward( timestep_r: torch.Tensor, encoder_hidden_states: torch.Tensor, context_latents: torch.Tensor, + joint_attention_kwargs: Optional[dict] = None, return_dict: bool = True, ) -> Union[torch.Tensor, Transformer2DModelOutput]: """The [`AceStepTransformer1DModel`] forward method. diff --git a/src/diffusers/pipelines/ace_step/pipeline_ace_step.py b/src/diffusers/pipelines/ace_step/pipeline_ace_step.py index 95800309ade2..e4b2979f2e37 100644 --- a/src/diffusers/pipelines/ace_step/pipeline_ace_step.py +++ b/src/diffusers/pipelines/ace_step/pipeline_ace_step.py @@ -225,6 +225,10 @@ def guidance_scale(self) -> float: def num_timesteps(self) -> int: return self._num_timesteps + @property + def joint_attention_kwargs(self): + return self._joint_attention_kwargs + def check_inputs( self, prompt: Union[str, List[str]], @@ -826,6 +830,7 @@ def __call__( cfg_interval_start: float = 0.0, cfg_interval_end: float = 1.0, timesteps: Optional[List[float]] = None, + joint_attention_kwargs: Optional[dict] = None, ): r""" The call function to the pipeline for music generation. @@ -984,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._joint_attention_kwargs = joint_attention_kwargs self._interrupt = False # Auto-generate instruction based on task_type if not provided @@ -1177,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), + joint_attention_kwargs=self.joint_attention_kwargs, return_dict=False, ) vt_cond, vt_uncond = model_output[0].chunk(2, dim=0) @@ -1200,6 +1207,7 @@ def __call__( timestep_r=t_curr_tensor, encoder_hidden_states=encoder_hidden_states, context_latents=context_latents, + joint_attention_kwargs=self.joint_attention_kwargs, return_dict=False, ) vt = model_output[0] @@ -1212,6 +1220,7 @@ def __call__( timestep_r=t_curr_tensor, encoder_hidden_states=non_cover_encoder_hidden_states, context_latents=context_latents, + joint_attention_kwargs=self.joint_attention_kwargs, return_dict=False, ) vt_nc = nc_output[0] diff --git a/tests/lora/test_lora_layers_ace_step.py b/tests/lora/test_lora_layers_ace_step.py index edb8ccf36608..8be764e368fe 100644 --- a/tests/lora/test_lora_layers_ace_step.py +++ b/tests/lora/test_lora_layers_ace_step.py @@ -189,22 +189,6 @@ def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): def test_correct_lora_configs_with_different_ranks(self): pass - @unittest.skip("AceStepPipeline does not accept attention_kwargs for LoRA scale.") - def test_simple_inference_with_text_denoiser_lora_and_scale(self): - pass - - @unittest.skip("AceStepPipeline does not accept attention_kwargs for LoRA scale.") - def test_lora_scale_kwargs_match_fusion(self): - pass - - @unittest.skip("AceStepPipeline does not accept attention_kwargs for LoRA scale.") - def test_set_adapters_match_attention_kwargs(self): - pass - - @unittest.skip("AceStepPipeline does not accept attention_kwargs for LoRA scale.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - @unittest.skip("AceStep attention layers have no bias; lora_bias is not applicable.") def test_lora_B_bias(self): pass From 968b8da6b96bbd654e2e2ceb9f98e574882888cc Mon Sep 17 00:00:00 2001 From: zcy Date: Thu, 16 Jul 2026 15:38:35 +0800 Subject: [PATCH 5/9] fix: override test_lora_fuse_nan instead of skipping it The base test hardcodes block attribute names (transformer_blocks, blocks, etc.) but AceStep uses 'layers'. Override with AceStep-specific block access path instead of skipping. Co-Authored-By: Claude Opus 4.6 (1M context) --- tests/lora/test_lora_layers_ace_step.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/tests/lora/test_lora_layers_ace_step.py b/tests/lora/test_lora_layers_ace_step.py index d97454894bac..edb8ccf36608 100644 --- a/tests/lora/test_lora_layers_ace_step.py +++ b/tests/lora/test_lora_layers_ace_step.py @@ -209,6 +209,23 @@ def test_simple_inference_with_text_lora_and_scale(self): def test_lora_B_bias(self): pass - @unittest.skip("AceStep uses 'layers' not 'transformer_blocks'/'blocks'; base test hardcodes block names.") def test_lora_fuse_nan(self): - pass + import numpy as np + + components, _, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to("cpu") + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) + + pipe.transformer.add_adapter(denoiser_lora_config, "adapter-1") + + with torch.no_grad(): + pipe.transformer.layers[0].self_attn.to_q.lora_A["adapter-1"].weight += float("inf") + + with self.assertRaises(ValueError): + pipe.fuse_lora(safe_fusing=True) + + pipe.fuse_lora(safe_fusing=False) + out = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(np.isnan(out).all()) From 2c49f32339fddf1268dd2c0bc7e0da3a40aee76b Mon Sep 17 00:00:00 2001 From: zcy Date: Thu, 16 Jul 2026 15:46:34 +0800 Subject: [PATCH 6/9] feat: add joint_attention_kwargs to AceStepPipeline for per-step LoRA scale - Add @apply_lora_scale("joint_attention_kwargs") decorator and joint_attention_kwargs param to AceStepTransformer1DModel.forward - Add joint_attention_kwargs param to AceStepPipeline.__call__ and wire it through to all three transformer calls - Un-skip 4 LoRA scale tests that now pass (39 passed, 13 skipped) Co-Authored-By: Claude Opus 4.6 (1M context) --- .../models/transformers/ace_step_transformer.py | 4 +++- .../pipelines/ace_step/pipeline_ace_step.py | 9 +++++++++ tests/lora/test_lora_layers_ace_step.py | 16 ---------------- 3 files changed, 12 insertions(+), 17 deletions(-) diff --git a/src/diffusers/models/transformers/ace_step_transformer.py b/src/diffusers/models/transformers/ace_step_transformer.py index 83e66e61b1de..6ba1f17738d6 100644 --- a/src/diffusers/models/transformers/ace_step_transformer.py +++ b/src/diffusers/models/transformers/ace_step_transformer.py @@ -22,7 +22,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin -from ...utils import logging +from ...utils import apply_lora_scale, logging from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import ( AttentionBackendName, @@ -529,6 +529,7 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("joint_attention_kwargs") def forward( self, hidden_states: torch.Tensor, @@ -536,6 +537,7 @@ def forward( timestep_r: torch.Tensor, encoder_hidden_states: torch.Tensor, context_latents: torch.Tensor, + joint_attention_kwargs: Optional[dict] = None, return_dict: bool = True, ) -> Union[torch.Tensor, Transformer2DModelOutput]: """The [`AceStepTransformer1DModel`] forward method. diff --git a/src/diffusers/pipelines/ace_step/pipeline_ace_step.py b/src/diffusers/pipelines/ace_step/pipeline_ace_step.py index 95800309ade2..e4b2979f2e37 100644 --- a/src/diffusers/pipelines/ace_step/pipeline_ace_step.py +++ b/src/diffusers/pipelines/ace_step/pipeline_ace_step.py @@ -225,6 +225,10 @@ def guidance_scale(self) -> float: def num_timesteps(self) -> int: return self._num_timesteps + @property + def joint_attention_kwargs(self): + return self._joint_attention_kwargs + def check_inputs( self, prompt: Union[str, List[str]], @@ -826,6 +830,7 @@ def __call__( cfg_interval_start: float = 0.0, cfg_interval_end: float = 1.0, timesteps: Optional[List[float]] = None, + joint_attention_kwargs: Optional[dict] = None, ): r""" The call function to the pipeline for music generation. @@ -984,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._joint_attention_kwargs = joint_attention_kwargs self._interrupt = False # Auto-generate instruction based on task_type if not provided @@ -1177,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), + joint_attention_kwargs=self.joint_attention_kwargs, return_dict=False, ) vt_cond, vt_uncond = model_output[0].chunk(2, dim=0) @@ -1200,6 +1207,7 @@ def __call__( timestep_r=t_curr_tensor, encoder_hidden_states=encoder_hidden_states, context_latents=context_latents, + joint_attention_kwargs=self.joint_attention_kwargs, return_dict=False, ) vt = model_output[0] @@ -1212,6 +1220,7 @@ def __call__( timestep_r=t_curr_tensor, encoder_hidden_states=non_cover_encoder_hidden_states, context_latents=context_latents, + joint_attention_kwargs=self.joint_attention_kwargs, return_dict=False, ) vt_nc = nc_output[0] diff --git a/tests/lora/test_lora_layers_ace_step.py b/tests/lora/test_lora_layers_ace_step.py index edb8ccf36608..8be764e368fe 100644 --- a/tests/lora/test_lora_layers_ace_step.py +++ b/tests/lora/test_lora_layers_ace_step.py @@ -189,22 +189,6 @@ def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): def test_correct_lora_configs_with_different_ranks(self): pass - @unittest.skip("AceStepPipeline does not accept attention_kwargs for LoRA scale.") - def test_simple_inference_with_text_denoiser_lora_and_scale(self): - pass - - @unittest.skip("AceStepPipeline does not accept attention_kwargs for LoRA scale.") - def test_lora_scale_kwargs_match_fusion(self): - pass - - @unittest.skip("AceStepPipeline does not accept attention_kwargs for LoRA scale.") - def test_set_adapters_match_attention_kwargs(self): - pass - - @unittest.skip("AceStepPipeline does not accept attention_kwargs for LoRA scale.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - @unittest.skip("AceStep attention layers have no bias; lora_bias is not applicable.") def test_lora_B_bias(self): pass From 2cd20398862d869fe7e645f48d30785c425d9687 Mon Sep 17 00:00:00 2001 From: zcy Date: Thu, 16 Jul 2026 18:07:37 +0800 Subject: [PATCH 7/9] refactor: use attention_kwargs instead of joint_attention_kwargs, use torch_device in test - Rename joint_attention_kwargs to attention_kwargs in transformer forward and pipeline __call__ for consistency with the diffusers naming convention for new pipelines - Add docstring for the attention_kwargs parameter - Use torch_device instead of hardcoded "cpu" in test_lora_fuse_nan Co-Authored-By: Claude Opus 4.6 (1M context) --- .../models/transformers/ace_step_transformer.py | 7 +++++-- .../pipelines/ace_step/pipeline_ace_step.py | 14 +++++++------- tests/lora/test_lora_layers_ace_step.py | 3 ++- 3 files changed, 14 insertions(+), 10 deletions(-) diff --git a/src/diffusers/models/transformers/ace_step_transformer.py b/src/diffusers/models/transformers/ace_step_transformer.py index 6ba1f17738d6..821c7ad1491a 100644 --- a/src/diffusers/models/transformers/ace_step_transformer.py +++ b/src/diffusers/models/transformers/ace_step_transformer.py @@ -529,7 +529,7 @@ def __init__( self.gradient_checkpointing = False - @apply_lora_scale("joint_attention_kwargs") + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, @@ -537,7 +537,7 @@ def forward( timestep_r: torch.Tensor, encoder_hidden_states: torch.Tensor, context_latents: torch.Tensor, - joint_attention_kwargs: Optional[dict] = None, + attention_kwargs: Optional[dict] = None, return_dict: bool = True, ) -> Union[torch.Tensor, Transformer2DModelOutput]: """The [`AceStepTransformer1DModel`] forward method. @@ -554,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. diff --git a/src/diffusers/pipelines/ace_step/pipeline_ace_step.py b/src/diffusers/pipelines/ace_step/pipeline_ace_step.py index e4b2979f2e37..e9bcb14e22ef 100644 --- a/src/diffusers/pipelines/ace_step/pipeline_ace_step.py +++ b/src/diffusers/pipelines/ace_step/pipeline_ace_step.py @@ -226,8 +226,8 @@ def num_timesteps(self) -> int: return self._num_timesteps @property - def joint_attention_kwargs(self): - return self._joint_attention_kwargs + def attention_kwargs(self): + return self._attention_kwargs def check_inputs( self, @@ -830,7 +830,7 @@ def __call__( cfg_interval_start: float = 0.0, cfg_interval_end: float = 1.0, timesteps: Optional[List[float]] = None, - joint_attention_kwargs: Optional[dict] = None, + attention_kwargs: Optional[dict] = None, ): r""" The call function to the pipeline for music generation. @@ -989,7 +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._joint_attention_kwargs = joint_attention_kwargs + self._attention_kwargs = attention_kwargs self._interrupt = False # Auto-generate instruction based on task_type if not provided @@ -1183,7 +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), - joint_attention_kwargs=self.joint_attention_kwargs, + attention_kwargs=self.attention_kwargs, return_dict=False, ) vt_cond, vt_uncond = model_output[0].chunk(2, dim=0) @@ -1207,7 +1207,7 @@ def __call__( timestep_r=t_curr_tensor, encoder_hidden_states=encoder_hidden_states, context_latents=context_latents, - joint_attention_kwargs=self.joint_attention_kwargs, + attention_kwargs=self.attention_kwargs, return_dict=False, ) vt = model_output[0] @@ -1220,7 +1220,7 @@ def __call__( timestep_r=t_curr_tensor, encoder_hidden_states=non_cover_encoder_hidden_states, context_latents=context_latents, - joint_attention_kwargs=self.joint_attention_kwargs, + attention_kwargs=self.attention_kwargs, return_dict=False, ) vt_nc = nc_output[0] diff --git a/tests/lora/test_lora_layers_ace_step.py b/tests/lora/test_lora_layers_ace_step.py index 8be764e368fe..990077fadfac 100644 --- a/tests/lora/test_lora_layers_ace_step.py +++ b/tests/lora/test_lora_layers_ace_step.py @@ -26,6 +26,7 @@ from ..testing_utils import ( require_peft_backend, skip_mps, + torch_device, ) @@ -198,7 +199,7 @@ def test_lora_fuse_nan(self): components, _, denoiser_lora_config = self.get_dummy_components() pipe = self.pipeline_class(**components) - pipe = pipe.to("cpu") + pipe = pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) From 09e2255e0ca3bc9c63d0022bfa11f23a0c397d3c Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Thu, 16 Jul 2026 11:33:35 +0000 Subject: [PATCH 8/9] Apply style fixes --- src/diffusers/loaders/lora_conversion_utils.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index 2c22a06f991b..07e3351685e8 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -3099,11 +3099,9 @@ def convert_module(module): 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 + 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 = { From 944e7590850805a4f6beb9b1bc45e26f58db2a7a Mon Sep 17 00:00:00 2001 From: chenyangzhu1 Date: Thu, 16 Jul 2026 19:47:45 +0800 Subject: [PATCH 9/9] docs: document ACE-Step attention kwargs --- src/diffusers/pipelines/ace_step/pipeline_ace_step.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/diffusers/pipelines/ace_step/pipeline_ace_step.py b/src/diffusers/pipelines/ace_step/pipeline_ace_step.py index e9bcb14e22ef..b11e08208eb1 100644 --- a/src/diffusers/pipelines/ace_step/pipeline_ace_step.py +++ b/src/diffusers/pipelines/ace_step/pipeline_ace_step.py @@ -914,6 +914,9 @@ def __call__( End ratio (0.0-1.0) of the timestep range where CFG is applied. timesteps (`List[float]`, *optional*): Custom timestep schedule. If provided, overrides `num_inference_steps` and `shift`. + attention_kwargs (`dict`, *optional*): + A kwargs dictionary passed along to the `AttentionProcessor`. Used to pass the LoRA scale via + `{"scale": float}`. Examples: