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

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>

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

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,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,
+            sampled_cfg_type,
-        omega_scale,
+            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,
+        cfg_type=sampled_cfg_type,
-        omega_scale=omega_scale,
+        omega_scale=sampled_omega_scale,
         manual_seeds=manual_seeds,
         guidance_interval=guidance_interval,
         guidance_interval_decay=guidance_interval_decay,

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

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

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": [
+        # Only needed to train or fine-tune; inference does not import these.
-            "peft",
+        "train": read_requirements("requirements-train.txt"),
-            "tensorboard",
-            "tensorboardX"
-        ]
     },
 )