Merge pull request #136 from ace-step/add_lora_support

add chinese_rap_lora
This commit is contained in:
Gong Junmin
2025-05-13 15:01:17 +08:00
committed by GitHub
56 changed files with 1333 additions and 177 deletions
+14 -5
View File
@@ -3,19 +3,28 @@ from pathlib import Path
import random
DEFAULT_ROOT_DIR = "examples/input_params"
DEFAULT_ROOT_DIR = "examples/default/input_params"
ZH_RAP_LORA_ROOT_DIR = "examples/zh_rap_lora/input_params"
class DataSampler:
def __init__(self, root_dir=DEFAULT_ROOT_DIR):
self.root_dir = root_dir
self.input_params_files = list(Path(self.root_dir).glob("*.json"))
self.zh_rap_lora_input_params_files = list(Path(ZH_RAP_LORA_ROOT_DIR).glob("*.json"))
self.zh_rap_lora_input_params_files += list(Path(ZH_RAP_LORA_ROOT_DIR).glob("*.json"))
def load_json(self, file_path):
with open(file_path, "r", encoding="utf-8") as f:
return json.load(f)
def sample(self):
json_path = random.choice(self.input_params_files)
json_data = self.load_json(json_path)
def sample(self, lora_name_or_path=None):
if lora_name_or_path is None or lora_name_or_path == "none":
json_path = random.choice(self.input_params_files)
json_data = self.load_json(json_path)
else:
json_path = random.choice(self.zh_rap_lora_input_params_files)
json_data = self.load_json(json_path)
# Update the lora_name in the json_data
json_data["lora_name_or_path"] = lora_name_or_path
return json_data
+1
View File
@@ -71,6 +71,7 @@ def main(checkpoint_path, server_name, port, device_id, share, bf16, torch_compi
demo = create_main_demo_ui(
text2music_process_func=model_demo.__call__,
sample_data_func=data_sampler.sample,
load_data_func=data_sampler.load_json,
)
demo.launch(server_name=server_name, server_port=port, share=share)
+22 -2
View File
@@ -113,6 +113,7 @@ class ACEStepPipeline:
ensure_directory_exists(checkpoint_dir)
self.checkpoint_dir = checkpoint_dir
self.lora_path = "none"
device = (
torch.device(f"cuda:{device_id}")
if torch.cuda.is_available()
@@ -1563,6 +1564,21 @@ class ACEStepPipeline:
input_audio = input_audio.to(device=device, dtype=dtype)
latents, _ = self.music_dcae.encode(input_audio, sr=sr)
return latents
def load_lora(self, lora_name_or_path):
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)
else:
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)
logger.info(f"Loading lora weights from: {lora_name_or_path} download path is: {lora_download_path}")
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.")
self.ace_step_transformer.unload_lora()
def __call__(
self,
@@ -1587,6 +1603,7 @@ class ACEStepPipeline:
audio2audio_enable: bool = False,
ref_audio_strength: float = 0.5,
ref_audio_input: str = None,
lora_name_or_path: str = "none",
retake_seeds: list = None,
retake_variance: float = 0.5,
task: str = "text2music",
@@ -1615,8 +1632,10 @@ class ACEStepPipeline:
self.load_quantized_checkpoint(self.checkpoint_dir)
else:
self.load_checkpoint(self.checkpoint_dir)
load_model_cost = time.time() - start_time
logger.info(f"Model loaded in {load_model_cost:.2f} seconds.")
self.load_lora(lora_name_or_path)
load_model_cost = time.time() - start_time
logger.info(f"Model loaded in {load_model_cost:.2f} seconds.")
start_time = time.time()
@@ -1813,6 +1832,7 @@ class ACEStepPipeline:
}
input_params_json = {
"lora_name_or_path": lora_name_or_path,
"task": task,
"prompt": prompt if task != "edit" else edit_target_prompt,
"lyrics": lyrics if task != "edit" else edit_target_lyrics,
+122 -22
View File
@@ -8,6 +8,7 @@ Apache 2.0 License
import gradio as gr
import librosa
import os
TAG_DEFAULT = "funk, pop, soul, rock, melodic, guitar, drums, bass, keyboard, percussion, 105 BPM, energetic, upbeat, groovy, vibrant, dynamic"
@@ -48,6 +49,26 @@ Catch the tune and hold it tight
In this moment we take flight
"""
# First, let's define the presets at the top of the file, after the imports
GENRE_PRESETS = {
"Modern Pop": "pop, synth, drums, guitar, 120 bpm, upbeat, catchy, vibrant, female vocals, polished vocals",
"Rock": "rock, electric guitar, drums, bass, 130 bpm, energetic, rebellious, gritty, male vocals, raw vocals",
"Hip Hop": "hip hop, 808 bass, hi-hats, synth, 90 bpm, bold, urban, intense, male vocals, rhythmic vocals",
"Country": "country, acoustic guitar, steel guitar, fiddle, 100 bpm, heartfelt, rustic, warm, male vocals, twangy vocals",
"EDM": "edm, synth, bass, kick drum, 128 bpm, euphoric, pulsating, energetic, instrumental",
"Reggae": "reggae, guitar, bass, drums, 80 bpm, chill, soulful, positive, male vocals, smooth vocals",
"Classical": "classical, orchestral, strings, piano, 60 bpm, elegant, emotive, timeless, instrumental",
"Jazz": "jazz, saxophone, piano, double bass, 110 bpm, smooth, improvisational, soulful, male vocals, crooning vocals",
"Metal": "metal, electric guitar, double kick drum, bass, 160 bpm, aggressive, intense, heavy, male vocals, screamed vocals",
"R&B": "r&b, synth, bass, drums, 85 bpm, sultry, groovy, romantic, female vocals, silky vocals"
}
# Add this function to handle preset selection
def update_tags_from_preset(preset_name):
if preset_name == "Custom":
return ""
return GENRE_PRESETS.get(preset_name, "")
def create_output_ui(task_name="Text2Music"):
# For many consumer-grade GPU devices, only one batch can be run
@@ -69,7 +90,17 @@ def create_text2music_ui(
gr,
text2music_process_func,
sample_data_func=None,
load_data_func=None,
):
with gr.Row(equal_height=True):
curr_file_dir = os.path.dirname(__file__)
output_file_dir = os.path.join(curr_file_dir, "..", "..", "outputs")
json_files = [f for f in os.listdir(output_file_dir) if f.endswith('.json')]
json_files.sort(reverse=True, key=lambda x: int(x.split('_')[1]))
output_files = gr.Dropdown(choices=json_files, label="Select previous generated input params", scale=9, interactive=True)
load_bnt = gr.Button("Load", variant="primary", scale=1)
with gr.Row():
with gr.Column():
with gr.Row(equal_height=True):
@@ -84,10 +115,18 @@ def create_text2music_ui(
info="-1 means random duration (30 ~ 240).",
scale=9,
)
sample_bnt = gr.Button("Sample", variant="primary", scale=1)
sample_bnt = gr.Button("Sample", variant="secondary", scale=1)
# audio2audio
audio2audio_enable = gr.Checkbox(label="Enable Audio2Audio", value=False, info="Check to enable Audio-to-Audio generation using a reference audio.", elem_id="audio2audio_checkbox")
with gr.Row(equal_height=True):
audio2audio_enable = gr.Checkbox(label="Enable Audio2Audio", value=False, info="Check to enable Audio-to-Audio generation using a reference audio.", elem_id="audio2audio_checkbox")
lora_name_or_path = gr.Dropdown(
label="Lora Name or Path",
choices=["ACE-Step/ACE-Step-v1-chinese-rap-LoRA", "none"],
value="none",
allow_custom_value=True,
)
ref_audio_input = gr.Audio(type="filepath", label="Reference Audio (for Audio2Audio)", visible=False, elem_id="ref_audio_input", show_download_button=True)
ref_audio_strength = gr.Slider(
label="Refer audio strength",
@@ -112,33 +151,51 @@ def create_text2music_ui(
outputs=[ref_audio_input, ref_audio_strength],
)
prompt = gr.Textbox(
lines=2,
label="Tags",
max_lines=4,
value=TAG_DEFAULT,
info="Support tags, descriptions, and scene. Use commas to separate different tags.\ntags and lyrics examples are from ai music generation community",
)
lyrics = gr.Textbox(
lines=9,
label="Lyrics",
max_lines=13,
value=LYRIC_DEFAULT,
info="Support lyric structure tags like [verse], [chorus], and [bridge] to separate different parts of the lyrics.\nUse [instrumental] or [inst] to generate instrumental music. Not support genre structure tag in lyrics",
with gr.Column(scale=2):
with gr.Group():
gr.Markdown("""<center>Support tags, descriptions, and scene. Use commas to separate different tags.<br>Tags and lyrics examples are from AI music generation community.</center>""")
with gr.Row():
genre_preset = gr.Dropdown(
choices=["Custom"] + list(GENRE_PRESETS.keys()),
value="Custom",
label="Preset",
scale=1,
)
prompt = gr.Textbox(
lines=1,
label="Tags",
max_lines=4,
value=TAG_DEFAULT,
scale=9,
)
# Add the change event for the preset dropdown
genre_preset.change(
fn=update_tags_from_preset,
inputs=[genre_preset],
outputs=[prompt]
)
with gr.Group():
gr.Markdown("""<center>Support lyric structure tags like [verse], [chorus], and [bridge] to separate different parts of the lyrics.<br>Use [instrumental] or [inst] to generate instrumental music. Not support genre structure tag in lyrics</center>""")
lyrics = gr.Textbox(
lines=9,
label="Lyrics",
max_lines=13,
value=LYRIC_DEFAULT,
)
with gr.Accordion("Basic Settings", open=False):
infer_step = gr.Slider(
minimum=1,
maximum=1000,
maximum=200,
step=1,
value=27,
value=60,
label="Infer Steps",
interactive=True,
)
guidance_scale = gr.Slider(
minimum=0.0,
maximum=200.0,
maximum=30.0,
step=0.1,
value=15.0,
label="Guidance Scale",
@@ -192,7 +249,7 @@ def create_text2music_ui(
)
use_erg_lyric = gr.Checkbox(
label="use ERG for lyric",
value=True,
value=False,
info="The same but apply to lyric encoder's attention.",
)
use_erg_diffusion = gr.Checkbox(
@@ -290,6 +347,7 @@ def create_text2music_ui(
retake_seeds=retake_seeds,
retake_variance=retake_variance,
task="retake",
lora_name_or_path="none" if "lora_name_or_path" not in json_data else json_data["lora_name_or_path"]
)
retake_bnt.click(
@@ -412,6 +470,7 @@ def create_text2music_ui(
repaint_start=repaint_start,
repaint_end=repaint_end,
src_audio_path=src_audio_path,
lora_name_or_path="none" if "lora_name_or_path" not in json_data else json_data["lora_name_or_path"]
)
repaint_bnt.click(
@@ -585,6 +644,7 @@ def create_text2music_ui(
edit_n_min=edit_n_min,
edit_n_max=edit_n_max,
retake_seeds=retake_seeds,
lora_name_or_path="none" if "lora_name_or_path" not in json_data else json_data["lora_name_or_path"]
)
edit_bnt.click(
@@ -729,6 +789,7 @@ def create_text2music_ui(
repaint_start=repaint_start,
repaint_end=repaint_end,
src_audio_path=src_audio_path,
lora_name_or_path="none" if "lora_name_or_path" not in json_data else json_data["lora_name_or_path"]
)
extend_bnt.click(
@@ -762,8 +823,7 @@ def create_text2music_ui(
outputs=extend_outputs + [extend_input_params_json],
)
def sample_data():
json_data = sample_data_func()
def json2output(json_data):
return (
json_data["audio_duration"],
json_data["prompt"],
@@ -808,8 +868,46 @@ def create_text2music_ui(
),
)
def sample_data(lora_name_or_path_):
json_data = sample_data_func(lora_name_or_path_)
return json2output(json_data)
sample_bnt.click(
sample_data,
inputs=[lora_name_or_path],
outputs=[
audio_duration,
prompt,
lyrics,
infer_step,
guidance_scale,
scheduler_type,
cfg_type,
omega_scale,
manual_seeds,
guidance_interval,
guidance_interval_decay,
min_guidance_scale,
use_erg_tag,
use_erg_lyric,
use_erg_diffusion,
oss_steps,
guidance_scale_text,
guidance_scale_lyric,
audio2audio_enable,
ref_audio_strength,
ref_audio_input,
],
)
def load_data(json_file):
json_file = os.path.join(output_file_dir, json_file)
json_data = load_data_func(json_file)
return json2output(json_data)
load_bnt.click(
fn=load_data,
inputs=[output_files],
outputs=[
audio_duration,
prompt,
@@ -859,6 +957,7 @@ def create_text2music_ui(
audio2audio_enable,
ref_audio_strength,
ref_audio_input,
lora_name_or_path,
],
outputs=outputs + [input_params_json],
)
@@ -867,6 +966,7 @@ def create_text2music_ui(
def create_main_demo_ui(
text2music_process_func=dump_func,
sample_data_func=dump_func,
load_data_func=dump_func,
):
with gr.Blocks(
title="ACE-Step Model 1.0 DEMO",
@@ -876,12 +976,12 @@ def create_main_demo_ui(
<h1 style="text-align: center;">ACE-Step: A Step Towards Music Generation Foundation Model</h1>
"""
)
with gr.Tab("text2music"):
create_text2music_ui(
gr=gr,
text2music_process_func=text2music_process_func,
sample_data_func=sample_data_func,
load_data_func=load_data_func,
)
return demo