From db5a9fea4ee166553fa19787ca720150ae833387 Mon Sep 17 00:00:00 2001 From: Andy <40263822+austin2035@users.noreply.github.com> Date: Thu, 15 May 2025 23:08:57 +0800 Subject: [PATCH 1/8] Update trainer.py fix lora trainer bug --- trainer.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/trainer.py b/trainer.py index b713d50..6aa3ee2 100644 --- a/trainer.py +++ b/trainer.py @@ -76,6 +76,7 @@ class Pipeline(LightningModule): except ImportError: raise ImportError("Please install peft library to use LoRA training") with open(lora_config_path, encoding="utf-8") as f: + import json lora_config = json.load(f) lora_config = LoraConfig(**lora_config) transformers.add_adapter(adapter_config=lora_config, adapter_name=adapter_name) @@ -825,6 +826,7 @@ def main(args): dataset_path=args.dataset_path, checkpoint_dir=args.checkpoint_dir, adapter_name=args.exp_name, + lora_config_path=args.lora_config_path ) checkpoint_callback = ModelCheckpoint( monitor=None, From 3d6fe064f5ec21ab60f900c7246e160f248d6796 Mon Sep 17 00:00:00 2001 From: woctordho Date: Sat, 17 May 2025 20:40:21 +0800 Subject: [PATCH 2/8] Enable gradient checkpointing --- acestep/models/ace_step_transformer.py | 20 ++------------------ acestep/models/lyrics_utils/lyric_encoder.py | 4 ++-- trainer.py | 1 + 3 files changed, 5 insertions(+), 20 deletions(-) diff --git a/acestep/models/ace_step_transformer.py b/acestep/models/ace_step_transformer.py index 391c985..3798042 100644 --- a/acestep/models/ace_step_transformer.py +++ b/acestep/models/ace_step_transformer.py @@ -360,10 +360,6 @@ class ACEStepTransformer2DModel( for module in self.children(): fn_recursive_feed_forward(module, chunk_size, dim) - def _set_gradient_checkpointing(self, module, value=False): - if hasattr(module, "gradient_checkpointing"): - module.gradient_checkpointing = value - def forward_lyric_encoder( self, lyric_token_idx: Optional[torch.LongTensor] = None, @@ -456,20 +452,8 @@ class ACEStepTransformer2DModel( if self.training and self.gradient_checkpointing: - def create_custom_forward(module, return_dict=None): - def custom_forward(*inputs): - if return_dict is not None: - return module(*inputs, return_dict=return_dict) - else: - return module(*inputs) - - return custom_forward - - ckpt_kwargs: Dict[str, Any] = ( - {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} - ) hidden_states = torch.utils.checkpoint.checkpoint( - create_custom_forward(block), + block, hidden_states=hidden_states, attention_mask=attention_mask, encoder_hidden_states=encoder_hidden_states, @@ -477,7 +461,7 @@ class ACEStepTransformer2DModel( rotary_freqs_cis=rotary_freqs_cis, rotary_freqs_cis_cross=encoder_rotary_freqs_cis, temb=temb, - **ckpt_kwargs, + use_reentrant=False, ) else: diff --git a/acestep/models/lyrics_utils/lyric_encoder.py b/acestep/models/lyrics_utils/lyric_encoder.py index bd849bd..6fe8f27 100644 --- a/acestep/models/lyrics_utils/lyric_encoder.py +++ b/acestep/models/lyrics_utils/lyric_encoder.py @@ -1030,8 +1030,8 @@ class ConformerEncoder(torch.nn.Module): mask_pad: torch.Tensor, ) -> torch.Tensor: for layer in self.encoders: - xs, chunk_masks, _, _ = ckpt.checkpoint( - layer.__call__, xs, chunk_masks, pos_emb, mask_pad + xs, chunk_masks, _, _ = torch.utils.checkpoint.checkpoint( + layer.__call__, xs, chunk_masks, pos_emb, mask_pad, use_reentrant=False ) return xs diff --git a/trainer.py b/trainer.py index 6aa3ee2..a47f4c4 100644 --- a/trainer.py +++ b/trainer.py @@ -68,6 +68,7 @@ class Pipeline(LightningModule): acestep_pipeline.load_checkpoint(acestep_pipeline.checkpoint_dir) transformers = acestep_pipeline.ace_step_transformer.float().cpu() + transformers.enable_gradient_checkpointing() assert lora_config_path is not None, "Please provide a LoRA config path" if lora_config_path is not None: From 2f6046b072b9ad8ad508bca435862697fe8a03a2 Mon Sep 17 00:00:00 2001 From: woctordho Date: Sat, 17 May 2025 22:58:28 +0800 Subject: [PATCH 3/8] Set scheduler interval to step and fewer warmup steps --- trainer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/trainer.py b/trainer.py index a47f4c4..a3ed9cb 100644 --- a/trainer.py +++ b/trainer.py @@ -49,7 +49,7 @@ class Pipeline(LightningModule): ssl_coeff: float = 1.0, checkpoint_dir=None, max_steps: int = 200000, - warmup_steps: int = 4000, + warmup_steps: int = 10, dataset_path: str = "./data/your_dataset_path", lora_config_path: str = None, adapter_name: str = "lora_adapter", @@ -440,7 +440,7 @@ class Pipeline(LightningModule): lr_scheduler = torch.optim.lr_scheduler.LambdaLR( optimizer, lr_lambda, last_epoch=-1 ) - return [optimizer], lr_scheduler + return [optimizer], [{"scheduler": lr_scheduler, "interval": "step"}] def train_dataloader(self): self.train_dataset = Text2MusicDataset( From c3abb452417e0fad2750699f408f5df4aeec8b60 Mon Sep 17 00:00:00 2001 From: Skelp Date: Sat, 17 May 2025 18:14:29 +0200 Subject: [PATCH 4/8] #194: implement negative tags for text2audio --- acestep/pipeline_ace_step.py | 68 +++++++++++++++++++++++++++++++----- acestep/ui/components.py | 10 ++++++ 2 files changed, 69 insertions(+), 9 deletions(-) diff --git a/acestep/pipeline_ace_step.py b/acestep/pipeline_ace_step.py index 6a41c89..0893425 100644 --- a/acestep/pipeline_ace_step.py +++ b/acestep/pipeline_ace_step.py @@ -1027,6 +1027,8 @@ class ACEStepPipeline: min_guidance_scale=3.0, oss_steps=[], encoder_text_hidden_states_null=None, + neg_encoder_text_hidden_states=None, + neg_text_attention_mask=None, use_erg_lyric=False, use_erg_diffusion=False, retake_random_generators=None, @@ -1323,14 +1325,24 @@ class ACEStepPipeline: }, ) else: - # P(null_speaker, null_text, null_lyric) - encoder_hidden_states_null, _ = self.ace_step_transformer.encode( - torch.zeros_like(encoder_text_hidden_states), - text_attention_mask, - torch.zeros_like(speaker_embds), - torch.zeros_like(lyric_token_ids), - lyric_mask, - ) + # Using negative prompt for unconditional guidance + if neg_encoder_text_hidden_states is not None: + encoder_hidden_states_null, _ = self.ace_step_transformer.encode( + neg_encoder_text_hidden_states, # Already padded to match + neg_text_attention_mask, # Already padded to match + torch.zeros_like(speaker_embds), + torch.zeros_like(lyric_token_ids), + lyric_mask, + ) + else: + # Original approach with zeros + encoder_hidden_states_null, _ = self.ace_step_transformer.encode( + torch.zeros_like(encoder_text_hidden_states), + text_attention_mask, + torch.zeros_like(speaker_embds), + torch.zeros_like(lyric_token_ids), + lyric_mask, + ) encoder_hidden_states_no_lyric = None if do_double_condition_guidance: @@ -1631,6 +1643,7 @@ class ACEStepPipeline: format: str = "wav", audio_duration: float = 60.0, prompt: str = None, + negative_prompt: str = None, lyrics: str = None, infer_step: int = 60, guidance_scale: float = 15.0, @@ -1702,6 +1715,40 @@ class ACEStepPipeline: encoder_text_hidden_states = encoder_text_hidden_states.repeat(batch_size, 1, 1) text_attention_mask = text_attention_mask.repeat(batch_size, 1) + if negative_prompt: + neg_texts = [negative_prompt] + neg_encoder_text_hidden_states, neg_text_attention_mask = self.get_text_embeddings( + neg_texts, self.device + ) + neg_encoder_text_hidden_states = neg_encoder_text_hidden_states.repeat(batch_size, 1, 1) + neg_text_attention_mask = neg_text_attention_mask.repeat(batch_size, 1) + + # Determine which is longer and pad the shorter one + pos_seq_len = encoder_text_hidden_states.shape[1] + neg_seq_len = neg_encoder_text_hidden_states.shape[1] + + if pos_seq_len > neg_seq_len: + # Pad negative embeddings + pad_size = pos_seq_len - neg_seq_len + neg_encoder_text_hidden_states = torch.nn.functional.pad( + neg_encoder_text_hidden_states, (0, 0, 0, pad_size), "constant", 0 + ) + neg_text_attention_mask = torch.nn.functional.pad( + neg_text_attention_mask, (0, pad_size), "constant", 0 + ) + elif neg_seq_len > pos_seq_len: + # Pad positive embeddings + pad_size = neg_seq_len - pos_seq_len + encoder_text_hidden_states = torch.nn.functional.pad( + encoder_text_hidden_states, (0, 0, 0, pad_size), "constant", 0 + ) + text_attention_mask = torch.nn.functional.pad( + text_attention_mask, (0, pad_size), "constant", 0 + ) + else: + neg_encoder_text_hidden_states = None + neg_text_attention_mask = None + encoder_text_hidden_states_null = None if use_erg_tag: encoder_text_hidden_states_null = self.get_text_embeddings_null( @@ -1843,7 +1890,9 @@ class ACEStepPipeline: guidance_interval_decay=guidance_interval_decay, min_guidance_scale=min_guidance_scale, oss_steps=oss_steps, - encoder_text_hidden_states_null=encoder_text_hidden_states_null, + encoder_text_hidden_states_null=encoder_text_hidden_states_null, + neg_encoder_text_hidden_states=neg_encoder_text_hidden_states if negative_prompt else None, + neg_text_attention_mask=neg_text_attention_mask if negative_prompt else None, use_erg_lyric=use_erg_lyric, use_erg_diffusion=use_erg_diffusion, retake_random_generators=retake_random_generators, @@ -1883,6 +1932,7 @@ class ACEStepPipeline: "lora_name_or_path": lora_name_or_path, "task": task, "prompt": prompt if task != "edit" else edit_target_prompt, + "negative_prompt": negative_prompt, "lyrics": lyrics if task != "edit" else edit_target_lyrics, "audio_duration": audio_duration, "infer_step": infer_step, diff --git a/acestep/ui/components.py b/acestep/ui/components.py index 3b4980e..f699ff7 100644 --- a/acestep/ui/components.py +++ b/acestep/ui/components.py @@ -174,6 +174,15 @@ def create_text2music_ui( scale=9, ) + with gr.Row(): + negative_prompt = gr.Textbox( + lines=1, + label="Negative Tags", + info="Elements you want to avoid in the generated music (only effective when no EFG is used)", + max_lines=4, + scale=9, + ) + # Add the change event for the preset dropdown genre_preset.change( fn=update_tags_from_preset, @@ -949,6 +958,7 @@ def create_text2music_ui( format, audio_duration, prompt, + negative_prompt, lyrics, infer_step, guidance_scale, From 508fc67bad550347ac201bb2d99e9cd52fb7f951 Mon Sep 17 00:00:00 2001 From: Roberts Slisans Date: Sat, 17 May 2025 20:55:43 +0300 Subject: [PATCH 5/8] split out training packages --- setup.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/setup.py b/setup.py index 89ade60..a2696d7 100644 --- a/setup.py +++ b/setup.py @@ -24,4 +24,11 @@ setup( package_data={ "acestep.models.lyrics_utils": ["vocab.json"], # Specify the relative path to vocab.json }, + extras_require={ + "train": [ + "peft", + "tensorboard", + "tensorboardX" + ] + }, ) From 7453602789aa5511d4bbce30810116269458b5e2 Mon Sep 17 00:00:00 2001 From: Roberts Slisans Date: Sun, 18 May 2025 13:39:52 +0300 Subject: [PATCH 6/8] bump version due to new features when installing the package, not updating the version number causes pip not to install the update --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index a2696d7..b58461b 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( description="ACE Step: A Step Towards Music Generation Foundation Model", long_description=open("README.md", encoding="utf-8").read(), long_description_content_type="text/markdown", - version="0.1.2", + version="0.2.0", packages=find_namespace_packages(), install_requires=open("requirements.txt", encoding="utf-8").read().splitlines(), author="ACE Studio, StepFun AI", From e3a6f4f659cdb073e415c811e1bb9d3c0de567f3 Mon Sep 17 00:00:00 2001 From: Gong Junmin <1836678486@qq.com> Date: Sun, 18 May 2025 19:42:29 +0800 Subject: [PATCH 7/8] Revert "#194: implement negative tags for text2audio" --- acestep/pipeline_ace_step.py | 68 +++++------------------------------- acestep/ui/components.py | 10 ------ 2 files changed, 9 insertions(+), 69 deletions(-) diff --git a/acestep/pipeline_ace_step.py b/acestep/pipeline_ace_step.py index 0893425..6a41c89 100644 --- a/acestep/pipeline_ace_step.py +++ b/acestep/pipeline_ace_step.py @@ -1027,8 +1027,6 @@ class ACEStepPipeline: min_guidance_scale=3.0, oss_steps=[], encoder_text_hidden_states_null=None, - neg_encoder_text_hidden_states=None, - neg_text_attention_mask=None, use_erg_lyric=False, use_erg_diffusion=False, retake_random_generators=None, @@ -1325,24 +1323,14 @@ class ACEStepPipeline: }, ) else: - # Using negative prompt for unconditional guidance - if neg_encoder_text_hidden_states is not None: - encoder_hidden_states_null, _ = self.ace_step_transformer.encode( - neg_encoder_text_hidden_states, # Already padded to match - neg_text_attention_mask, # Already padded to match - torch.zeros_like(speaker_embds), - torch.zeros_like(lyric_token_ids), - lyric_mask, - ) - else: - # Original approach with zeros - encoder_hidden_states_null, _ = self.ace_step_transformer.encode( - torch.zeros_like(encoder_text_hidden_states), - text_attention_mask, - torch.zeros_like(speaker_embds), - torch.zeros_like(lyric_token_ids), - lyric_mask, - ) + # P(null_speaker, null_text, null_lyric) + encoder_hidden_states_null, _ = self.ace_step_transformer.encode( + torch.zeros_like(encoder_text_hidden_states), + text_attention_mask, + torch.zeros_like(speaker_embds), + torch.zeros_like(lyric_token_ids), + lyric_mask, + ) encoder_hidden_states_no_lyric = None if do_double_condition_guidance: @@ -1643,7 +1631,6 @@ class ACEStepPipeline: format: str = "wav", audio_duration: float = 60.0, prompt: str = None, - negative_prompt: str = None, lyrics: str = None, infer_step: int = 60, guidance_scale: float = 15.0, @@ -1715,40 +1702,6 @@ class ACEStepPipeline: encoder_text_hidden_states = encoder_text_hidden_states.repeat(batch_size, 1, 1) text_attention_mask = text_attention_mask.repeat(batch_size, 1) - if negative_prompt: - neg_texts = [negative_prompt] - neg_encoder_text_hidden_states, neg_text_attention_mask = self.get_text_embeddings( - neg_texts, self.device - ) - neg_encoder_text_hidden_states = neg_encoder_text_hidden_states.repeat(batch_size, 1, 1) - neg_text_attention_mask = neg_text_attention_mask.repeat(batch_size, 1) - - # Determine which is longer and pad the shorter one - pos_seq_len = encoder_text_hidden_states.shape[1] - neg_seq_len = neg_encoder_text_hidden_states.shape[1] - - if pos_seq_len > neg_seq_len: - # Pad negative embeddings - pad_size = pos_seq_len - neg_seq_len - neg_encoder_text_hidden_states = torch.nn.functional.pad( - neg_encoder_text_hidden_states, (0, 0, 0, pad_size), "constant", 0 - ) - neg_text_attention_mask = torch.nn.functional.pad( - neg_text_attention_mask, (0, pad_size), "constant", 0 - ) - elif neg_seq_len > pos_seq_len: - # Pad positive embeddings - pad_size = neg_seq_len - pos_seq_len - encoder_text_hidden_states = torch.nn.functional.pad( - encoder_text_hidden_states, (0, 0, 0, pad_size), "constant", 0 - ) - text_attention_mask = torch.nn.functional.pad( - text_attention_mask, (0, pad_size), "constant", 0 - ) - else: - neg_encoder_text_hidden_states = None - neg_text_attention_mask = None - encoder_text_hidden_states_null = None if use_erg_tag: encoder_text_hidden_states_null = self.get_text_embeddings_null( @@ -1890,9 +1843,7 @@ class ACEStepPipeline: guidance_interval_decay=guidance_interval_decay, min_guidance_scale=min_guidance_scale, oss_steps=oss_steps, - encoder_text_hidden_states_null=encoder_text_hidden_states_null, - neg_encoder_text_hidden_states=neg_encoder_text_hidden_states if negative_prompt else None, - neg_text_attention_mask=neg_text_attention_mask if negative_prompt else None, + encoder_text_hidden_states_null=encoder_text_hidden_states_null, use_erg_lyric=use_erg_lyric, use_erg_diffusion=use_erg_diffusion, retake_random_generators=retake_random_generators, @@ -1932,7 +1883,6 @@ class ACEStepPipeline: "lora_name_or_path": lora_name_or_path, "task": task, "prompt": prompt if task != "edit" else edit_target_prompt, - "negative_prompt": negative_prompt, "lyrics": lyrics if task != "edit" else edit_target_lyrics, "audio_duration": audio_duration, "infer_step": infer_step, diff --git a/acestep/ui/components.py b/acestep/ui/components.py index f699ff7..3b4980e 100644 --- a/acestep/ui/components.py +++ b/acestep/ui/components.py @@ -174,15 +174,6 @@ def create_text2music_ui( scale=9, ) - with gr.Row(): - negative_prompt = gr.Textbox( - lines=1, - label="Negative Tags", - info="Elements you want to avoid in the generated music (only effective when no EFG is used)", - max_lines=4, - scale=9, - ) - # Add the change event for the preset dropdown genre_preset.change( fn=update_tags_from_preset, @@ -958,7 +949,6 @@ def create_text2music_ui( format, audio_duration, prompt, - negative_prompt, lyrics, infer_step, guidance_scale, From 375431486c8533b221fed64ca79972b1c0fc9354 Mon Sep 17 00:00:00 2001 From: GeorgeDrb Date: Sun, 18 May 2025 15:35:21 +0300 Subject: [PATCH 8/8] Add memory cleanup to prevent VRAM overflow --- acestep/pipeline_ace_step.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/acestep/pipeline_ace_step.py b/acestep/pipeline_ace_step.py index 0893425..6567050 100644 --- a/acestep/pipeline_ace_step.py +++ b/acestep/pipeline_ace_step.py @@ -135,6 +135,21 @@ class ACEStepPipeline: self.cpu_offload = cpu_offload self.quantized = quantized self.overlapped_decode = overlapped_decode + + def cleanup_memory(self): + """Clean up GPU and CPU memory to prevent VRAM overflow during multiple generations.""" + # Clear CUDA cache + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + # Log memory usage if in verbose mode + allocated = torch.cuda.memory_allocated() / (1024 ** 3) + reserved = torch.cuda.memory_reserved() / (1024 ** 3) + logger.info(f"GPU Memory: {allocated:.2f}GB allocated, {reserved:.2f}GB reserved") + + # Collect Python garbage + import gc + gc.collect() def load_checkpoint(self, checkpoint_dir=None, export_quantized_weights=False): device = self.device @@ -1918,6 +1933,9 @@ class ACEStepPipeline: save_path=save_path, format=format, ) + + # Clean up memory after generation + self.cleanup_memory() end_time = time.time() latent2audio_time_cost = end_time - start_time