Add Song Describer pipeline: prepare_song_describer.py, prepare_from_hf, preprocess/train CLI

acestep/training_v2/cli/args.py

@@ -108,6 +108,75 @@ def build_root_parser() -> argparse.ArgumentParser:
         help="Random seed (default: 42)",
     )
 
+    # -- from-hf ------------------------------------------------------------
+    p_from_hf = subparsers.add_parser(
+        "from-hf",
+        help="Prepare ACE-Step dataset from a Hugging Face dataset (writes dataset.json + audio)",
+        formatter_class=formatter_class,
+    )
+    p_from_hf.add_argument(
+        "--dataset",
+        type=str,
+        required=True,
+        metavar="NAME",
+        help="Hugging Face dataset id (e.g. ashrafemam/crema-d or polyai/minds14)",
+    )
+    p_from_hf.add_argument(
+        "--output-dir",
+        type=str,
+        required=True,
+        metavar="DIR",
+        help="Output directory for dataset.json and audio/",
+    )
+    p_from_hf.add_argument(
+        "--split",
+        type=str,
+        default="train",
+        help="Dataset split (default: train)",
+    )
+    p_from_hf.add_argument(
+        "--config",
+        type=str,
+        default=None,
+        help="Dataset config name if required",
+    )
+    p_from_hf.add_argument(
+        "--caption-column",
+        type=str,
+        default=None,
+        help="Column to use as caption (default: auto-detect caption/text/sentence)",
+    )
+    p_from_hf.add_argument(
+        "--audio-column",
+        type=str,
+        default=None,
+        help="Column containing audio (default: auto-detect)",
+    )
+    p_from_hf.add_argument(
+        "--max-samples",
+        type=int,
+        default=None,
+        help="Max number of samples to export (default: all)",
+    )
+    p_from_hf.add_argument(
+        "--audio-subdir",
+        type=str,
+        default="audio",
+        help="Subdirectory name for audio under output-dir (default: audio)",
+    )
+    p_from_hf.add_argument(
+        "--json-filename",
+        type=str,
+        default="dataset.json",
+        help="Output JSON filename (default: dataset.json)",
+    )
+    p_from_hf.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        default=False,
+        help="Allow loading datasets with custom code",
+    )
+
     return root
 
 

acestep/training_v2/prepare_from_hf.py

@@ -0,0 +1,194 @@
+"""
+Prepare an ACE-Step–compatible dataset from a Hugging Face dataset.
+
+Loads a HF dataset (with an audio column and optional text/caption column),
+writes audio files to a local directory and a dataset JSON in the format
+expected by ``preprocess_audio_files``.
+
+Usage (standalone, only needs ``pip install datasets``):
+    python prepare_from_hf_cli.py --dataset <HF_ID> --output-dir <DIR> [--max-samples N]
+
+Or via train.py (full env):
+    python train.py from-hf --dataset <HF_ID> --output-dir <DIR>
+
+Then preprocess and train:
+    python train.py preprocess --dataset-json <out>/dataset.json --tensor-output <pt_dir> ...
+    python train.py fixed --dataset-dir <pt_dir> ...
+
+Datasets with an "audio" column (HF Audio feature) are supported; each row
+must provide either a path or decoded bytes. Caption is taken from a
+configurable column (default: caption/text/sentence). Note: google/MusicCaps
+on HF has no audio column (YouTube refs only); use a dataset that includes
+audio (e.g. polyai/minds14, ashrafemam/crema-d) or add audio separately.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import shutil
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+# Default column names to use as caption (first present wins)
+DEFAULT_CAPTION_COLUMNS = ("caption", "text", "sentence", "description", "transcript")
+
+
+def _infer_audio_column(column_names: List[str], first_row: Dict[str, Any]) -> Optional[str]:
+    for c in column_names:
+        if c in first_row and first_row[c] is not None:
+            val = first_row[c]
+            if isinstance(val, dict) and ("path" in val or "bytes" in val):
+                return c
+            if isinstance(val, str) and Path(val).suffix.lower() in {".wav", ".mp3", ".flac", ".ogg", ".m4a", ".opus"}:
+                return c
+    return None
+
+
+def _infer_caption_column(column_names: List[str], first_row: Dict[str, Any]) -> Optional[str]:
+    for name in DEFAULT_CAPTION_COLUMNS:
+        if name in column_names and first_row.get(name) and isinstance(first_row[name], str):
+            return name
+    return None
+
+
+def _audio_path_from_row(audio_val: Any, audio_dir: Path, index: int, suffix: str = ".wav") -> Optional[Path]:
+    if audio_val is None:
+        return None
+    if isinstance(audio_val, dict):
+        path = audio_val.get("path")
+        if path and Path(path).is_file():
+            dest = audio_dir / f"sample_{index:06d}{suffix}"
+            try:
+                shutil.copy2(path, dest)
+                return dest
+            except OSError as e:
+                logger.warning("Copy failed for row %d: %s", index, e)
+                return None
+        raw_bytes = audio_val.get("bytes")
+        if raw_bytes is not None:
+            dest = audio_dir / f"sample_{index:06d}{suffix}"
+            try:
+                dest.write_bytes(raw_bytes)
+                return dest
+            except OSError as e:
+                logger.warning("Write failed for row %d: %s", index, e)
+                return None
+        return None
+    if isinstance(audio_val, str) and Path(audio_val).is_file():
+        ext = Path(audio_val).suffix.lower() or suffix
+        dest = audio_dir / f"sample_{index:06d}{ext}"
+        try:
+            shutil.copy2(audio_val, dest)
+            return dest
+        except OSError as e:
+            logger.warning("Copy failed for row %d: %s", index, e)
+            return None
+    return None
+
+
+def prepare_from_hf(
+    dataset_name: str,
+    output_dir: str,
+    *,
+    split: str = "train",
+    config: Optional[str] = None,
+    caption_column: Optional[str] = None,
+    audio_column: Optional[str] = None,
+    max_samples: Optional[int] = None,
+    audio_subdir: str = "audio",
+    json_filename: str = "dataset.json",
+    trust_remote_code: bool = False,
+) -> Dict[str, Any]:
+    """Load a Hugging Face dataset and write ACE-Step dataset JSON + audio files.
+
+    Args:
+        dataset_name: Hugging Face dataset id (e.g. "google/MusicCaps" or "ashrafemam/crema-d").
+        output_dir: Directory to write dataset.json and audio files (into output_dir/<audio_subdir>).
+        split: Dataset split to use (default: "train").
+        config: Dataset config name if required.
+        caption_column: Column to use as caption; if None, inferred (caption/text/sentence/...).
+        audio_column: Column containing audio (path or Audio dict); if None, inferred.
+        max_samples: Limit number of samples (default: no limit).
+        audio_subdir: Subdirectory under output_dir for audio files (default: "audio").
+        json_filename: Name of the dataset JSON file (default: "dataset.json").
+        trust_remote_code: Passed to load_dataset.
+
+    Returns:
+        Dict with keys: output_dir, dataset_json, audio_dir, num_samples, caption_column, audio_column.
+    """
+    try:
+        from datasets import load_dataset
+    except ImportError:
+        raise ImportError("Install the 'datasets' package: pip install datasets") from None
+
+    out_path = Path(output_dir)
+    out_path.mkdir(parents=True, exist_ok=True)
+    audio_dir = out_path / audio_subdir
+    audio_dir.mkdir(parents=True, exist_ok=True)
+
+    load_kw: Dict[str, Any] = {"path": dataset_name, "split": split, "trust_remote_code": trust_remote_code}
+    if config:
+        load_kw["name"] = config
+    ds = load_dataset(**load_kw)
+    if hasattr(ds, "column_names"):
+        column_names = ds.column_names
+        first_row = ds[0] if len(ds) > 0 else {}
+    else:
+        column_names = list(ds[split].column_names)
+        first_row = ds[split][0] if len(ds[split]) > 0 else {}
+
+    audio_col = audio_column or _infer_audio_column(column_names, first_row)
+    if not audio_col:
+        raise ValueError(
+            "No audio column found. Ensure the dataset has an 'audio' column (Audio feature) "
+            "or pass --audio-column. For text-only datasets (e.g. MusicCaps with YouTube refs), "
+            "download audio separately and build the JSON manually."
+        )
+
+    caption_col = caption_column or _infer_caption_column(column_names, first_row)
+    data_split = ds[split] if hasattr(ds, "__getitem__") and split in ds else ds
+    total = len(data_split)
+    if max_samples is not None and max_samples > 0:
+        total = min(total, max_samples)
+
+    samples: List[Dict[str, Any]] = []
+    for i in range(total):
+        row = data_split[i]
+        audio_val = row.get(audio_col)
+        rel_audio_path = _audio_path_from_row(audio_val, audio_dir, i)
+        if rel_audio_path is None:
+            logger.debug("Skipping row %d: no resolvable audio", i)
+            continue
+        caption = (caption_col and row.get(caption_col)) or "[Instrumental]"
+        if not isinstance(caption, str):
+            caption = str(caption) if caption is not None else "[Instrumental]"
+        samples.append({
+            "filename": rel_audio_path.name,
+            "audio_path": str(rel_audio_path),
+            "caption": caption[:512],
+            "lyrics": "[Instrumental]",
+            "genre": "",
+            "bpm": None,
+            "keyscale": "",
+            "timesignature": "",
+            "duration": 0,
+            "is_instrumental": True,
+        })
+
+    dataset_json_path = out_path / json_filename
+    for s in samples:
+        s["audio_path"] = str(Path(audio_subdir) / Path(s["audio_path"]).name)
+    with open(dataset_json_path, "w", encoding="utf-8") as f:
+        json.dump({"samples": samples, "metadata": {"tag_position": "prepend", "genre_ratio": 0, "custom_tag": ""}}, f, indent=2)
+
+    return {
+        "output_dir": str(out_path),
+        "dataset_json": str(dataset_json_path),
+        "audio_dir": str(audio_dir),
+        "num_samples": len(samples),
+        "caption_column": caption_col,
+        "audio_column": audio_col,
+    }

prepare_from_hf_cli.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+"""
+Standalone CLI to prepare an ACE-Step dataset from a Hugging Face dataset.
+
+Only requires: pip install datasets
+
+Usage:
+    python prepare_from_hf_cli.py --dataset <HF_DATASET_ID> --output-dir <DIR> [options]
+
+Example:
+    python prepare_from_hf_cli.py --dataset polyai/minds14 --output-dir ./data/minds14 --split train
+
+Then preprocess and train:
+    python train.py preprocess --dataset-json ./data/minds14/dataset.json --tensor-output ./pt_minds14 ...
+    python train.py fixed --dataset-dir ./pt_minds14 ...
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description="Prepare ACE-Step dataset from a Hugging Face dataset (dataset.json + audio/)",
+    )
+    parser.add_argument("--dataset", required=True, metavar="NAME", help="Hugging Face dataset id")
+    parser.add_argument("--output-dir", required=True, metavar="DIR", help="Output directory for dataset.json and audio/")
+    parser.add_argument("--split", default="train", help="Dataset split (default: train)")
+    parser.add_argument("--config", default=None, help="Dataset config name if required")
+    parser.add_argument("--caption-column", default=None, help="Caption column (default: auto-detect)")
+    parser.add_argument("--audio-column", default=None, help="Audio column (default: auto-detect)")
+    parser.add_argument("--max-samples", type=int, default=None, help="Max samples to export (default: all)")
+    parser.add_argument("--audio-subdir", default="audio", help="Audio subdir under output-dir (default: audio)")
+    parser.add_argument("--json-filename", default="dataset.json", help="Output JSON filename (default: dataset.json)")
+    parser.add_argument("--trust-remote-code", action="store_true", help="Allow datasets with custom code")
+    args = parser.parse_args()
+
+    from acestep.training_v2.prepare_from_hf import prepare_from_hf
+
+    try:
+        result = prepare_from_hf(
+            dataset_name=args.dataset,
+            output_dir=args.output_dir,
+            split=args.split,
+            config=args.config,
+            caption_column=args.caption_column,
+            audio_column=args.audio_column,
+            max_samples=args.max_samples,
+            audio_subdir=args.audio_subdir,
+            json_filename=args.json_filename,
+            trust_remote_code=args.trust_remote_code,
+        )
+    except ImportError as e:
+        print(f"[FAIL] {e}", file=sys.stderr)
+        return 1
+    except Exception as e:
+        print(f"[FAIL] {e}", file=sys.stderr)
+        return 1
+
+    print(f"\n[OK] Prepared {result['num_samples']} samples")
+    print(f"     dataset_json: {result['dataset_json']}")
+    print(f"     audio_dir:    {result['audio_dir']}")
+    print("\nNext: preprocess then train (see train.py preprocess / train.py fixed).")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

prepare_song_describer.py

@@ -0,0 +1,248 @@
+#!/usr/bin/env python3
+"""
+Prepare ACE-Step dataset from audio.zip + song_describer.csv, then optionally preprocess and train.
+
+Downloads (or uses local) audio.zip and song_describer.csv, unzips audio, builds dataset.json
+in the format expected by train.py preprocess, then runs preprocess and train.
+
+Dataset source (Song Describer Dataset, SDD):
+    Zenodo: https://zenodo.org/records/10072001
+    - audio.zip (~3.3 GB, 706 recordings)
+    - song_describer.csv (~186 KB, ~1.1k captions)
+    Direct file URLs (use if the record page lists these names):
+    - https://zenodo.org/records/10072001/files/audio.zip
+    - https://zenodo.org/records/10072001/files/song_describer.csv
+
+Usage:
+    python prepare_song_describer.py --audio-zip <URL_or_path> --csv <URL_or_path> --output-dir <DIR> [options]
+
+Example (download from Zenodo then preprocess + train):
+    python prepare_song_describer.py --audio-zip "https://zenodo.org/records/10072001/files/audio.zip" --csv "https://zenodo.org/records/10072001/files/song_describer.csv" --output-dir ./data/song_describer --checkpoint-dir ./checkpoints --run-preprocess --run-train
+
+Example (local files):
+    python prepare_song_describer.py --audio-zip ./audio.zip --csv ./song_describer.csv --output-dir ./data/song_describer --checkpoint-dir ./checkpoints --run-preprocess --run-train
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import shutil
+import subprocess
+import sys
+import urllib.request
+from pathlib import Path
+
+
+AUDIO_SUBDIR = "audio"
+DATASET_JSON_NAME = "dataset.json"
+DEFAULT_CAPTION_COLUMNS = ("caption", "description", "text", "title", "label")
+DEFAULT_AUDIO_COLUMNS = ("filename", "path", "file", "audio", "id", "name")
+
+
+def _is_url(s: str) -> bool:
+    return s.strip().startswith(("http://", "https://"))
+
+
+def _download(url: str, dest: Path) -> None:
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    req = urllib.request.Request(url, headers={"User-Agent": "ACE-Step/1.0"})
+    with urllib.request.urlopen(req) as resp:
+        dest.write_bytes(resp.read())
+    print(f"[INFO] Downloaded {url} -> {dest}", file=sys.stderr)
+
+
+def _ensure_file(src: str, dest: Path) -> Path:
+    if _is_url(src):
+        _download(src, dest)
+        return dest
+    p = Path(src)
+    if not p.is_file():
+        raise FileNotFoundError(f"Not a file: {p}")
+    if p.resolve() != dest.resolve():
+        shutil.copy2(p, dest)
+    return dest
+
+
+def _infer_csv_columns(reader: csv.DictReader) -> tuple[str, str]:
+    names = [c for c in reader.fieldnames or [] if c]
+    if not names:
+        raise ValueError("CSV has no header columns")
+    audio_col = None
+    for c in DEFAULT_AUDIO_COLUMNS:
+        if c in names:
+            audio_col = c
+            break
+    if not audio_col:
+        audio_col = names[0]
+    caption_col = None
+    for c in DEFAULT_CAPTION_COLUMNS:
+        if c in names:
+            caption_col = c
+            break
+    if not caption_col:
+        caption_col = names[1] if len(names) > 1 else names[0]
+    return audio_col, caption_col
+
+
+def build_dataset_json(
+    csv_path: Path,
+    audio_dir: Path,
+    output_json_path: Path,
+    csv_audio_col: str | None = None,
+    csv_caption_col: str | None = None,
+) -> int:
+    audio_dir.mkdir(parents=True, exist_ok=True)
+    existing = {f.name for f in audio_dir.iterdir() if f.is_file()}
+
+    with open(csv_path, newline="", encoding="utf-8", errors="replace") as f:
+        reader = csv.DictReader(f)
+        audio_col, caption_col = _infer_csv_columns(reader)
+        if csv_audio_col:
+            audio_col = csv_audio_col
+        if csv_caption_col:
+            caption_col = csv_caption_col
+
+        samples: list[dict] = []
+        for row in reader:
+            raw_path = (row.get(audio_col) or "").strip()
+            if not raw_path:
+                continue
+            name = Path(raw_path).name
+            if name not in existing:
+                continue
+            caption = (row.get(caption_col) or "").strip() or "[Instrumental]"
+            if len(caption) > 512:
+                caption = caption[:512]
+            rel_audio = f"{AUDIO_SUBDIR}/{name}"
+            samples.append({
+                "filename": name,
+                "audio_path": rel_audio,
+                "caption": caption,
+                "lyrics": "[Instrumental]",
+                "genre": "",
+                "bpm": None,
+                "keyscale": "",
+                "timesignature": "",
+                "duration": 0,
+                "is_instrumental": True,
+            })
+
+    payload = {
+        "samples": samples,
+        "metadata": {"tag_position": "prepend", "genre_ratio": 0, "custom_tag": ""},
+    }
+    output_json_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(output_json_path, "w", encoding="utf-8") as out:
+        json.dump(payload, out, indent=2)
+    print(f"[INFO] Wrote {len(samples)} samples to {output_json_path}", file=sys.stderr)
+    return len(samples)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description="Prepare dataset from audio.zip + song_describer.csv, then preprocess/train."
+    )
+    parser.add_argument("--audio-zip", required=True, help="URL or path to audio.zip")
+    parser.add_argument("--csv", required=True, help="URL or path to song_describer.csv")
+    parser.add_argument("--output-dir", required=True, help="Output directory (dataset.json + audio/)")
+    parser.add_argument("--csv-audio-col", default=None, help="CSV column for audio filename (default: auto)")
+    parser.add_argument("--csv-caption-col", default=None, help="CSV column for caption (default: auto)")
+    parser.add_argument("--checkpoint-dir", default=None, help="Checkpoint dir for preprocess/train")
+    parser.add_argument("--run-preprocess", action="store_true", help="Run train.py preprocess after preparing")
+    parser.add_argument("--run-train", action="store_true", help="Run train.py fixed after preprocessing")
+    parser.add_argument("--tensor-output", default=None, help="Dir for .pt tensors (default: <output-dir>/tensors)")
+    parser.add_argument("--lora-output", default=None, help="Dir for LoRA output (default: <output-dir>/lora_output)")
+    parser.add_argument("--model-variant", default="turbo", help="Model variant (default: turbo)")
+    args = parser.parse_args()
+
+    out_dir = Path(args.output_dir).resolve()
+    work = out_dir / "work"
+    work.mkdir(parents=True, exist_ok=True)
+
+    zip_path = work / "audio.zip"
+    csv_path = work / "song_describer.csv"
+    _ensure_file(args.audio_zip, zip_path)
+    _ensure_file(args.csv, csv_path)
+
+    audio_dir = out_dir / AUDIO_SUBDIR
+    if zip_path.is_file():
+        tmp_extract = work / "audio_extract"
+        tmp_extract.mkdir(parents=True, exist_ok=True)
+        shutil.unpack_archive(str(zip_path), str(tmp_extract))
+        audio_dir.mkdir(parents=True, exist_ok=True)
+        for f in tmp_extract.rglob("*"):
+            if f.is_file():
+                dest = audio_dir / f.name
+                if dest != f.resolve():
+                    shutil.copy2(f, dest)
+        shutil.rmtree(tmp_extract, ignore_errors=True)
+        print(f"[INFO] Unpacked {zip_path} -> {audio_dir} (flattened)", file=sys.stderr)
+    else:
+        audio_dir.mkdir(parents=True, exist_ok=True)
+
+    dataset_json = out_dir / DATASET_JSON_NAME
+    n = build_dataset_json(
+        csv_path,
+        audio_dir,
+        dataset_json,
+        csv_audio_col=args.csv_audio_col,
+        csv_caption_col=args.csv_caption_col,
+    )
+    if n == 0:
+        print("[FAIL] No samples in dataset (CSV rows must match filenames in zip).", file=sys.stderr)
+        return 1
+
+    tensor_output = args.tensor_output or str(out_dir / "tensors")
+    lora_output = args.lora_output or str(out_dir / "lora_output")
+
+    if args.run_preprocess or args.run_train:
+        if not args.checkpoint_dir:
+            print("[FAIL] --checkpoint-dir required for --run-preprocess / --run-train.", file=sys.stderr)
+            return 1
+        train_py = Path(__file__).resolve().parent / "train.py"
+        if not train_py.is_file():
+            print(f"[FAIL] train.py not found: {train_py}", file=sys.stderr)
+            return 1
+
+    if args.run_preprocess:
+        cmd = [
+            sys.executable,
+            str(train_py),
+            "preprocess",
+            "--dataset-json", str(dataset_json),
+            "--tensor-output", tensor_output,
+            "--checkpoint-dir", args.checkpoint_dir,
+            "--model-variant", args.model_variant,
+        ]
+        print(f"[INFO] Running: {' '.join(cmd)}", file=sys.stderr)
+        rc = subprocess.call(cmd, cwd=str(train_py.parent))
+        if rc != 0:
+            return rc
+
+    if args.run_train:
+        cmd = [
+            sys.executable,
+            str(train_py),
+            "fixed",
+            "--dataset-dir", tensor_output,
+            "--checkpoint-dir", args.checkpoint_dir,
+            "--model-variant", args.model_variant,
+            "--output-dir", lora_output,
+        ]
+        print(f"[INFO] Running: {' '.join(cmd)}", file=sys.stderr)
+        rc = subprocess.call(cmd, cwd=str(train_py.parent))
+        if rc != 0:
+            return rc
+
+    print(f"\n[OK] Dataset: {dataset_json} ({n} samples)")
+    if not args.run_preprocess:
+        print(f"  Preprocess: python train.py preprocess --dataset-json {dataset_json} --tensor-output {tensor_output} --checkpoint-dir <ckpt> --model-variant {args.model_variant}")
+    if args.run_preprocess and not args.run_train:
+        print(f"  Train:     python train.py fixed --dataset-dir {tensor_output} --checkpoint-dir <ckpt> --model-variant {args.model_variant} --output-dir {lora_output}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

train.py

@@ -61,7 +61,7 @@ def _has_subcommand() -> bool:
     args = sys.argv[1:]
     if "--help" in args or "-h" in args:
         return True  # let argparse handle help
-    known = {"vanilla", "fixed", "estimate"}
+    known = {"vanilla", "fixed", "estimate", "from-hf"}
     return bool(known & set(args))
 
 
@@ -89,7 +89,10 @@ def _dispatch(args) -> int:
 
     sub = args.subcommand
 
-    # All subcommands need path validation
+    if sub == "from-hf":
+        return _run_from_hf(args)
+
+    # All other subcommands need path validation
     if not validate_paths(args):
         return 1
 
@@ -202,6 +205,52 @@ def _run_preprocess(args) -> int:
     return 0
 
 
+def _run_from_hf(args) -> int:
+    """Prepare dataset from a Hugging Face dataset."""
+    from acestep.training_v2.prepare_from_hf import prepare_from_hf
+
+    out_dir = getattr(args, "output_dir", None)
+    if not out_dir:
+        print("[FAIL] from-hf requires --output-dir.", file=sys.stderr)
+        return 1
+
+    print("\n" + "=" * 60)
+    print("  Prepare from Hugging Face dataset")
+    print("=" * 60)
+    print(f"  Dataset:   {args.dataset}")
+    print(f"  Split:     {getattr(args, 'split', 'train')}")
+    print(f"  Output:    {out_dir}")
+    print("=" * 60)
+
+    try:
+        result = prepare_from_hf(
+            dataset_name=args.dataset,
+            output_dir=out_dir,
+            split=getattr(args, "split", "train"),
+            config=getattr(args, "config", None),
+            caption_column=getattr(args, "caption_column", None),
+            audio_column=getattr(args, "audio_column", None),
+            max_samples=getattr(args, "max_samples", None),
+            audio_subdir=getattr(args, "audio_subdir", "audio"),
+            json_filename=getattr(args, "json_filename", "dataset.json"),
+            trust_remote_code=getattr(args, "trust_remote_code", False),
+        )
+    except Exception as exc:
+        print(f"[FAIL] {exc}", file=sys.stderr)
+        logger.exception("from-hf error")
+        return 1
+
+    print(f"\n[OK] Prepared {result['num_samples']} samples:")
+    print(f"     dataset_json: {result['dataset_json']}")
+    print(f"     audio_dir:    {result['audio_dir']}")
+    print("\nNext steps:")
+    print(f"  1. Preprocess:  python train.py preprocess --dataset-json {result['dataset_json']} \\")
+    print(f"                    --tensor-output <pt_dir> --checkpoint-dir <ckpt> --model-variant turbo")
+    print(f"  2. Train:       python train.py fixed --dataset-dir <pt_dir> --checkpoint-dir <ckpt> \\")
+    print(f"                    --model-variant turbo --output-dir <lora_output>")
+    return 0
+
+
 def _run_estimate(args) -> int:
     """Run gradient sensitivity estimation."""
     import json as _json