Spaces:
Sleeping
Sleeping
File size: 3,162 Bytes
37a9a36 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | 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))
|