scripts/run_waveform_causal_continuation.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
from pathlib import Path

from music3lab.editing.waveform_causal_continuation_runner import (
    load_waveform_causal_continuation_config,
    run_waveform_causal_continuation,
)


def main() -> int:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)
    check = subparsers.add_parser("check")
    check.add_argument("--config", type=Path, required=True)
    run = subparsers.add_parser("run")
    run.add_argument("--config", type=Path, required=True)
    run.add_argument("--corpus-root", type=Path, required=True)
    run.add_argument("--output-root", type=Path, required=True)
    run.add_argument("--memory-profile", choices=("h100", "24gb"), default="h100")
    run.add_argument("--user-audio", type=Path)
    args = parser.parse_args()
    if args.command == "check":
        config = load_waveform_causal_continuation_config(args.config)
        result = {"status": "READY", "schema_version": config["schema_version"]}
        code = 0
    else:
        result = run_waveform_causal_continuation(
            config_path=args.config,
            corpus_root=args.corpus_root,
            output_root=args.output_root,
            memory_profile=args.memory_profile,
            user_audio=args.user_audio,
        )
        code = 0 if result["status"] == "FUNCTIONAL_PASS" else 1
    print(json.dumps(result, sort_keys=True, separators=(",", ":")))
    return code


if __name__ == "__main__":
    raise SystemExit(main())