| 29 | |
| 30 | |
| 31 | def parse_args() -> argparse.Namespace: |
| 32 | parser = argparse.ArgumentParser(description="Python reference ACE-Step warmbench.") |
| 33 | parser.add_argument("--family", default="ace_step") |
| 34 | parser.add_argument("--checkpoint-dir", default=str(DEFAULT_CHECKPOINT_DIR)) |
| 35 | parser.add_argument("--config-path", default="acestep-v15-turbo") |
| 36 | parser.add_argument("--lm-model-path", default="acestep-5Hz-lm-1.7B") |
| 37 | parser.add_argument("--lm-backend", choices=["pt", "vllm", "mlx"], default="pt") |
| 38 | parser.add_argument("--lm-dtype", choices=["auto", "float32", "float16", "bfloat16"], default="auto") |
| 39 | parser.add_argument("--backend", choices=["cpu", "cuda", "mps", "xpu"], default="cuda") |
| 40 | parser.add_argument("--device", type=int, default=0) |
| 41 | parser.add_argument("--threads", type=int, default=8) |
| 42 | parser.add_argument("--warmup", type=int, default=0) |
| 43 | parser.add_argument("--iterations", type=int, default=1) |
| 44 | parser.add_argument("--timing-file", default="") |
| 45 | parser.add_argument("--output-dir", default="") |
| 46 | parser.add_argument("--noise-file", default="") |
| 47 | parser.add_argument("--request-json", default="") |
| 48 | parser.add_argument("--request-sequence-json", default="") |
| 49 | return parser.parse_args() |
| 50 | |
| 51 | |
| 52 | def summarize_audio(audio: np.ndarray, sample_rate: int) -> dict[str, Any]: |