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..07e3351685e8 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -3094,3 +3094,31 @@ 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..821c7ad1491a 100644 --- a/src/diffusers/models/transformers/ace_step_transformer.py +++ b/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,6 +529,7 @@ def __init__( self.gradient_checkpointing = False + @apply_lora_scale("attention_kwargs") def forward( self, hidden_states: torch.Tensor, @@ -535,6 +537,7 @@ def forward( 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. diff --git a/src/diffusers/pipelines/ace_step/pipeline_ace_step.py b/src/diffusers/pipelines/ace_step/pipeline_ace_step.py index 26c14d8bfac7..b11e08208eb1 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. @@ -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. @@ -908,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: @@ -983,6 +992,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 +1186,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 +1210,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 +1223,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] 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..990077fadfac --- /dev/null +++ b/tests/lora/test_lora_layers_ace_step.py @@ -0,0 +1,216 @@ +# 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, + torch_device, +) + + +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, "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"] + 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("AceStep attention layers have no bias; lora_bias is not applicable.") + def test_lora_B_bias(self): + pass + + def test_lora_fuse_nan(self): + import numpy as np + + components, _, denoiser_lora_config = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + 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())