From e9ea6b9babede9f91631a2b1fcb826078adf4cc1 Mon Sep 17 00:00:00 2001 From: Leonid Pershin Date: Tue, 8 Sep 2026 20:31:07 +0300 Subject: [PATCH] feat: generation options for infer.py, and split out training deps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- TRAIN_INSTRUCTION.md | 10 +++ infer.py | 143 ++++++++++++++++++++++++++++++++--------- requirements-train.txt | 7 ++ requirements.txt | 5 -- setup.py | 20 ++++-- 5 files changed, 145 insertions(+), 40 deletions(-) create mode 100644 requirements-train.txt diff --git a/TRAIN_INSTRUCTION.md b/TRAIN_INSTRUCTION.md index b254560..c74e0bb 100644 --- a/TRAIN_INSTRUCTION.md +++ b/TRAIN_INSTRUCTION.md @@ -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 diff --git a/infer.py b/infer.py index 6ee2ca0..9a65e05 100644 --- a/infer.py +++ b/infer.py @@ -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,44 +117,70 @@ def main(checkpoint_path, bf16, torch_compile, cpu_offload, overlapped_decode, d cpu_offload=cpu_offload, overlapped_decode=overlapped_decode ) - print(model_demo) - data_sampler = DataSampler() + if lyrics_file is not None: + with open(lyrics_file, "r", encoding="utf-8") as f: + lyrics = f.read() - json_data = data_sampler.sample() - json_data = sample_data(json_data) - print(json_data) + if prompt is None: + # No prompt given: keep the original behaviour and use a random example. + data_sampler = DataSampler() + json_data = data_sampler.sample() + ( + audio_duration, + prompt, + sampled_lyrics, + infer_step, + sampled_guidance_scale, + scheduler_type, + sampled_cfg_type, + sampled_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, + ) = 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 - ( - 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, - ) = json_data + 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, diff --git a/requirements-train.txt b/requirements-train.txt new file mode 100644 index 0000000..4d095e7 --- /dev/null +++ b/requirements-train.txt @@ -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 diff --git a/requirements.txt b/requirements.txt index 8fae102..259974e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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 \ No newline at end of file diff --git a/setup.py b/setup.py index b58461b..8847c2a 100644 --- a/setup.py +++ b/setup.py @@ -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"), }, )