Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| from __future__ import annotations | |
| import argparse | |
| import asyncio | |
| import json | |
| import resource | |
| import sys | |
| import time | |
| from pathlib import Path | |
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) | |
| from openmusic_analysis.application import build_service | |
| async def main() -> None: | |
| parser = argparse.ArgumentParser( | |
| description="Measure model load and per-representation latency" | |
| ) | |
| parser.add_argument("audio", type=Path) | |
| parser.add_argument("--lyrics", type=Path) | |
| parser.add_argument( | |
| "--representations", | |
| nargs="+", | |
| default=None, | |
| help="Load and measure only these representations (emotion models are opt-in).", | |
| ) | |
| args = parser.parse_args() | |
| service = build_service() | |
| lyrics = args.lyrics.read_text(encoding="utf-8") if args.lyrics else None | |
| measurements = {} | |
| started = time.perf_counter() | |
| if args.representations is None: | |
| representations = ["audio.global", "audio.temporal"] | |
| if lyrics is not None: | |
| representations.append("lyrics.global") | |
| else: | |
| representations = list(dict.fromkeys(args.representations)) | |
| await service.load_models(representations) | |
| measurements["model_startup_seconds"] = time.perf_counter() - started | |
| for representation in representations: | |
| started = time.perf_counter() | |
| await service.analyze( | |
| args.audio, | |
| lyrics=lyrics, | |
| requested_representations=[representation], | |
| track_id=None, | |
| content_identity=None, | |
| ) | |
| measurements[f"{representation}_seconds"] = time.perf_counter() - started | |
| peak_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss | |
| # macOS reports bytes; Linux (including the Docker image) reports KiB. | |
| measurements["peak_process_rss_bytes"] = int( | |
| peak_rss if sys.platform == "darwin" else peak_rss * 1024 | |
| ) | |
| try: | |
| import torch | |
| if torch.cuda.is_available(): | |
| measurements["peak_cuda_memory_bytes"] = torch.cuda.max_memory_allocated() | |
| except ImportError: | |
| pass | |
| print(json.dumps(measurements, indent=2)) | |
| if __name__ == "__main__": | |
| asyncio.run(main()) | |