mrleo1nid/ACE-Step
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add lora weight

Browse Source
This commit is contained in:
rkfg
2025-05-20 08:59:00 +03:00
parent 516792aa82
commit fc984d81fa
2 changed files with 32 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
acestep/pipeline_ace_step.py
+13 -9
View File
@@ -33,6 +33,7 @@ from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import (
retrieve_timesteps,
)
from diffusers.utils.torch_utils import randn_tensor
from diffusers.utils.peft_utils import set_weights_and_activate_adapters
from transformers import UMT5EncoderModel, AutoTokenizer
from acestep.language_segmentation import LangSegment, language_filters
@@ -138,18 +139,18 @@ 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()
@@ -1410,8 +1411,8 @@ class ACEStepPipeline:
input_audio = input_audio.to(device=self.device, dtype=self.dtype)
latents, _ = self.music_dcae.encode(input_audio, sr=sr)
return latents
def load_lora(self, lora_name_or_path):
def load_lora(self, lora_name_or_path, lora_weight):
if lora_name_or_path != self.lora_path and lora_name_or_path != "none":
if not os.path.exists(lora_name_or_path):
lora_download_path = snapshot_download(lora_name_or_path, cache_dir=self.checkpoint_dir)
@@ -1419,8 +1420,9 @@ class ACEStepPipeline:
lora_download_path = lora_name_or_path
if self.lora_path != "none":
self.ace_step_transformer.unload_lora()
self.ace_step_transformer.load_lora_adapter(os.path.join(lora_download_path, "pytorch_lora_weights.safetensors"), adapter_name="zh_rap_lora", with_alpha=True, prefix=None)
self.ace_step_transformer.load_lora_adapter(os.path.join(lora_download_path, "pytorch_lora_weights.safetensors"), adapter_name="ace_step_lora", with_alpha=True, prefix=None)
logger.info(f"Loading lora weights from: {lora_name_or_path} download path is: {lora_download_path}")
set_weights_and_activate_adapters(self.ace_step_transformer,["ace_step_lora"], [lora_weight])
self.lora_path = lora_name_or_path
elif self.lora_path != "none" and lora_name_or_path == "none":
logger.info("No lora weights to load.")
@@ -1451,6 +1453,7 @@ class ACEStepPipeline:
ref_audio_strength: float = 0.5,
ref_audio_input: str = None,
lora_name_or_path: str = "none",
lora_weight: float = 1.0,
retake_seeds: list = None,
retake_variance: float = 0.5,
task: str = "text2music",
@@ -1478,8 +1481,8 @@ class ACEStepPipeline:
self.load_quantized_checkpoint(self.checkpoint_dir)
else:
self.load_checkpoint(self.checkpoint_dir)
self.load_lora(lora_name_or_path)
self.load_lora(lora_name_or_path, lora_weight)
load_model_cost = time.time() - start_time
logger.info(f"Model loaded in {load_model_cost:.2f} seconds.")
@@ -1663,7 +1666,7 @@ class ACEStepPipeline:
save_path=save_path,
format=format,
)
# Clean up memory after generation
self.cleanup_memory()
@@ -1678,6 +1681,7 @@ class ACEStepPipeline:
input_params_json = {
"format": format,
"lora_name_or_path": lora_name_or_path,
"lora_weight": lora_weight,
"task": task,
"prompt": prompt if task != "edit" else edit_target_prompt,
"lyrics": lyrics if task != "edit" else edit_target_lyrics,