()
| 376 | |
| 377 | |
| 378 | def main() -> int: |
| 379 | args = parse_args() |
| 380 | if args.family != "stable_audio": |
| 381 | raise RuntimeError(f"unsupported Stable Audio warmbench family: {args.family}") |
| 382 | |
| 383 | os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") |
| 384 | os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") |
| 385 | os.environ.setdefault("STABLE_AUDIO_LOCAL_ONLY", "1") |
| 386 | os.environ["STABLE_AUDIO_MODEL_ROOT"] = str(resolve_repo_path(args.model_root).resolve()) |
| 387 | add_reference_path(args.reference_root) |
| 388 | device = normalize_device(args) |
| 389 | seed_all(args.seed) |
| 390 | |
| 391 | from stable_audio_3 import StableAudioModel |
| 392 | |
| 393 | model_name = model_name_from_arg(args.model) |
| 394 | model_half = args.model_precision == "fp16" or (args.model_precision == "native" and device == "cuda") |
| 395 | model = StableAudioModel.from_pretrained(model_name, device=device, model_half=model_half) |
| 396 | sync_device(device) |
| 397 | |
| 398 | requests = load_requests(args) |
| 399 | if not requests: |
| 400 | raise RuntimeError("Stable Audio warmbench request sequence is empty") |
| 401 | |
| 402 | output_root = args.output_dir if args.output_dir is not None else args.audio_out.parent |
| 403 | output_root = resolve_repo_path(output_root).resolve() |
| 404 | timing_path = resolve_repo_path(args.timing_file) |
| 405 | timing_path.parent.mkdir(parents=True, exist_ok=True) |
| 406 | timing_lines: list[str] = [] |
| 407 | |
| 408 | previous_audio: list[Path] = [] |
| 409 | for warmup_index in range(max(0, args.warmup)): |
| 410 | _, warmup_timing = run_request( |
| 411 | model, |
| 412 | requests[0], |
| 413 | output_root / "warmup" / f"{warmup_index:02d}", |
| 414 | args, |
| 415 | device, |
| 416 | previous_audio, |
| 417 | ) |
| 418 | timing_lines.extend(warmup_timing) |
| 419 | |
| 420 | steps: list[dict[str, Any]] = [] |
| 421 | for request_index, request in enumerate(requests): |
| 422 | total_ms = 0.0 |
| 423 | last_step: dict[str, Any] | None = None |
| 424 | for iteration in range(max(1, args.iterations)): |
| 425 | step, run_timing = run_request( |
| 426 | model, |
| 427 | request, |
| 428 | output_root / f"request_{request_index:02d}" / f"iter_{iteration:02d}", |
| 429 | args, |
| 430 | device, |
| 431 | previous_audio, |
| 432 | ) |
| 433 | total_ms += float(step["metrics"]["wall_ms"]) |
| 434 | last_step = step |
| 435 | timing_lines.extend(run_timing) |
no test coverage detected