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

feat: generation options for infer.py, and split out training deps

Browse Source 
infer.py could only render a random example from examples/input_params:
there was no way to pass your own prompt, lyrics, duration or seed, so
using it for anything specific meant writing a separate script. Add
--prompt, --lyrics/--lyrics_file, --duration, --steps, --guidance_scale,
--scheduler, --cfg_type, --omega_scale, --seed and --format alongside the
existing runtime flags. Without --prompt the old random-example behaviour
is kept, so existing invocations are unaffected.

Also move the training-only packages out of the default install.
datasets, pytorch_lightning, matplotlib, tensorboard and tensorboardX are
imported by trainer.py and convert2hf_dataset.py, never on the inference
path, but every user was installing them — and datasets==3.4.1 is a hard
pin that drags constraints onto huggingface-hub. They now live in
requirements-train.txt behind the existing (previously ineffective)
"train" extra: pip install -e ".[train]".

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Leonid Pershin
2026-09-08 20:31:07 +03:00
co-authored by Claude Opus 5
parent c7953dc4e0
commit e9ea6b9bab
5 changed files with 145 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
TRAIN_INSTRUCTION.md
+10
View File
@@ -1,5 +1,15 @@
# Training Instruction
## 0. Install the Training Dependencies
Training needs a few packages that a plain inference install does not pull in
(`datasets`, `pytorch_lightning`, `matplotlib`, `tensorboard`, `tensorboardX`).
Install them with the `train` extra:
```bash
pip install -e ".[train]"
```
## 1. Data Preparation
### Required File Format
infer.py
+99 -14
View File
@@ -48,7 +48,66 @@ def sample_data(json_data):
)
@click.option("--device_id", type=int, default=0, help="Device ID to use")
@click.option("--output_path", type=str, default=None, help="Path to save the output")
def main(checkpoint_path, bf16, torch_compile, cpu_offload, overlapped_decode, device_id, output_path):
# --- generation parameters -------------------------------------------------
# Without --prompt the script keeps its original behaviour and generates from a
# random example in examples/input_params.
@click.option(
"--prompt",
type=str,
default=None,
help="Style tags or description, comma separated. If omitted, a random example is used.",
)
@click.option("--lyrics", type=str, default=None, help="Lyrics, with [verse]/[chorus] tags.")
@click.option(
"--lyrics_file",
type=click.Path(exists=True, dir_okay=False),
default=None,
help="Read lyrics from a UTF-8 file. Takes precedence over --lyrics.",
)
@click.option("--duration", type=float, default=60.0, help="Audio duration in seconds (-1 for random).")
@click.option("--steps", type=int, default=60, help="Number of inference steps.")
@click.option("--guidance_scale", type=float, default=15.0, help="Guidance scale.")
@click.option(
"--scheduler",
type=click.Choice(["euler", "heun", "pingpong"]),
default="euler",
help="Scheduler type.",
)
@click.option(
"--cfg_type",
type=click.Choice(["cfg", "apg", "cfg_star"]),
default="apg",
help="CFG type. apg is recommended.",
)
@click.option("--omega_scale", type=float, default=10.0, help="Granularity scale.")
@click.option("--seed", type=str, default=None, help="Manual seeds, comma separated. Random if omitted.")
@click.option(
"--format",
"audio_format",
type=click.Choice(["wav", "mp3", "ogg", "flac"]),
default="wav",
help="Output audio format.",
)
def main(
checkpoint_path,
bf16,
torch_compile,
cpu_offload,
overlapped_decode,
device_id,
output_path,
prompt,
lyrics,
lyrics_file,
duration,
steps,
guidance_scale,
scheduler,
cfg_type,
omega_scale,
seed,
audio_format,
):
os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)
model_demo = ACEStepPipeline(
@@ -58,23 +117,24 @@ def main(checkpoint_path, bf16, torch_compile, cpu_offload, overlapped_decode, d
cpu_offload=cpu_offload,
overlapped_decode=overlapped_decode
)
print(model_demo)
if lyrics_file is not None:
with open(lyrics_file, "r", encoding="utf-8") as f:
lyrics = f.read()
if prompt is None:
# No prompt given: keep the original behaviour and use a random example.
data_sampler = DataSampler()
json_data = data_sampler.sample()
json_data = sample_data(json_data)
print(json_data)
(
audio_duration,
prompt,
lyrics,
sampled_lyrics,
infer_step,
guidance_scale,
sampled_guidance_scale,
scheduler_type,
cfg_type,
omega_scale,
sampled_cfg_type,
sampled_omega_scale,
manual_seeds,
guidance_interval,
guidance_interval_decay,
@@ -85,17 +145,42 @@ def main(checkpoint_path, bf16, torch_compile, cpu_offload, overlapped_decode, d
oss_steps,
guidance_scale_text,
guidance_scale_lyric,
) = json_data
) = sample_data(json_data)
# Explicit options still win over the sampled example.
lyrics = lyrics if lyrics is not None else sampled_lyrics
manual_seeds = seed if seed is not None else manual_seeds
else:
audio_duration = duration
lyrics = lyrics if lyrics is not None else "[instrumental]"
infer_step = steps
sampled_guidance_scale = guidance_scale
scheduler_type = scheduler
sampled_cfg_type = cfg_type
sampled_omega_scale = omega_scale
manual_seeds = seed
guidance_interval = 0.5
guidance_interval_decay = 0.0
min_guidance_scale = 3.0
use_erg_tag = True
use_erg_lyric = True
use_erg_diffusion = True
oss_steps = None
guidance_scale_text = 0.0
guidance_scale_lyric = 0.0
click.echo(f"prompt: {prompt}")
click.echo(f"duration: {audio_duration}s, steps: {infer_step}, seeds: {manual_seeds}")
model_demo(
format=audio_format,
audio_duration=audio_duration,
prompt=prompt,
lyrics=lyrics,
infer_step=infer_step,
guidance_scale=guidance_scale,
guidance_scale=sampled_guidance_scale,
scheduler_type=scheduler_type,
cfg_type=cfg_type,
omega_scale=omega_scale,
cfg_type=sampled_cfg_type,
omega_scale=sampled_omega_scale,
manual_seeds=manual_seeds,
guidance_interval=guidance_interval,
guidance_interval_decay=guidance_interval_decay,
requirements-train.txt
+7
View File
@@ -0,0 +1,7 @@
# Training-only dependencies (trainer.py, convert2hf_dataset.py).
# Install with: pip install -e ".[train]"
datasets==3.4.1
pytorch_lightning==2.5.1
matplotlib==3.10.1
tensorboard
tensorboardX
requirements.txt
-5
View File
@@ -1,12 +1,9 @@
datasets==3.4.1
diffusers>=0.33.0
gradio>=6.0.0
librosa==0.11.0
loguru==0.7.3
matplotlib==3.10.1
numpy
pypinyin==0.53.0
pytorch_lightning==2.5.1
soundfile==0.13.1
torch
torchaudio
@@ -22,5 +19,3 @@ cutlet
fugashi[unidic-lite]
click
peft
tensorboard
tensorboardX
setup.py
+14 -6
View File
@@ -1,5 +1,16 @@
from setuptools import setup, find_namespace_packages
def read_requirements(path):
"""Read a requirements file, skipping comments and blank lines."""
with open(path, encoding="utf-8") as f:
return [
line.strip()
for line in f
if line.strip() and not line.lstrip().startswith("#")
]
setup(
name="ace_step",
description="ACE Step: A Step Towards Music Generation Foundation Model",
@@ -7,7 +18,7 @@ setup(
long_description_content_type="text/markdown",
version="0.2.0",
packages=find_namespace_packages(),
install_requires=open("requirements.txt", encoding="utf-8").read().splitlines(),
install_requires=read_requirements("requirements.txt"),
author="ACE Studio, StepFun AI",
license="Apache 2.0",
classifiers=[
@@ -25,10 +36,7 @@ setup(
"acestep.models.lyrics_utils": ["vocab.json"], # Specify the relative path to vocab.json
},
extras_require={
"train": [
"peft",
"tensorboard",
"tensorboardX"
]
# Only needed to train or fine-tune; inference does not import these.
"train": read_requirements("requirements-train.txt"),
},
)