Spaces:
Sleeping
Sleeping
| import argparse | |
| import asyncio | |
| import logging | |
| import time | |
| import psutil | |
| import os | |
| import numpy as np | |
| from openmusic_analysis.settings import Settings | |
| from openmusic_analysis.worker.service import WorkerService | |
| logging.basicConfig(level=logging.INFO) | |
| log = logging.getLogger(__name__) | |
| async def run_benchmark(audio_file: str, iterations: int): | |
| # Enable unpinned models so benchmark can run with main revision locally | |
| os.environ["OWLSONIC_ALLOW_UNPINNED_MODELS"] = "true" | |
| settings = Settings.from_env() | |
| service = WorkerService(settings) | |
| log.info("Starting benchmark...") | |
| t0 = time.perf_counter() | |
| await service.startup() | |
| startup_ms = (time.perf_counter() - t0) * 1000 | |
| log.info(f"Startup complete in {startup_ms:.2f} ms") | |
| process = psutil.Process() | |
| rss_mb = process.memory_info().rss / (1024 * 1024) | |
| log.info(f"Memory RSS after startup: {rss_mb:.2f} MB") | |
| for i in range(iterations): | |
| log.info(f"--- Iteration {i+1}/{iterations} ---") | |
| try: | |
| t0 = time.perf_counter() | |
| embed_res = await service.audio_embeddings(audio_file) | |
| full_ms = (time.perf_counter() - t0) * 1000 | |
| log.info(f"Embeddings API time: {full_ms:.2f} ms") | |
| log.info(f" Probe time: {embed_res.timings.probe_ms:.2f} ms") | |
| log.info(f" Decode time: {embed_res.timings.decode_ms:.2f} ms") | |
| log.info(f" Inference time: {embed_res.timings.model_inference_ms:.2f} ms") | |
| # Pass to emotions | |
| t0 = time.perf_counter() | |
| temporal_dicts = None | |
| if embed_res.audio_temporal and embed_res.temporal_positions: | |
| temporal_dicts = [] | |
| for t_emb, pos in zip(embed_res.audio_temporal, embed_res.temporal_positions): | |
| temporal_dicts.append({ | |
| "start_ms": pos["start_ms"], | |
| "end_ms": pos["end_ms"], | |
| "embedding": t_emb | |
| }) | |
| emo_res = await service.emotions_from_audio_embeddings( | |
| embedding_version=embed_res.metadata.audio_embedding_version, | |
| global_embedding=embed_res.audio_global, | |
| temporal_embeddings=temporal_dicts | |
| ) | |
| emo_ms = (time.perf_counter() - t0) * 1000 | |
| log.info(f"Emotions API time: {emo_ms:.2f} ms") | |
| log.info(f" Inference time: {emo_res.timings.model_inference_ms:.2f} ms") | |
| except Exception as e: | |
| log.error(f"Benchmark iteration {i+1} failed: {e}") | |
| if (i + 1) % 10 == 0 or i == iterations - 1: | |
| rss_mb = process.memory_info().rss / (1024 * 1024) | |
| log.info(f"Memory RSS at iteration {i+1}: {rss_mb:.2f} MB") | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("audio", help="Path to audio file") | |
| parser.add_argument("--iters", type=int, default=1, help="Number of iterations") | |
| args = parser.parse_args() | |
| asyncio.run(run_benchmark(args.audio, args.iters)) | |