/
/
/
1"""Smart Fades audio analysis provider."""
2
3from __future__ import annotations
4
5import asyncio
6import time
7from dataclasses import dataclass, field
8from datetime import timedelta
9from typing import TYPE_CHECKING, Any
10
11import numpy as np
12import soxr
13import torch
14from beat_this.inference import Spect2Frames, aggregate_prediction, split_piece
15from music_assistant_models.config_entries import ConfigEntry
16from music_assistant_models.enums import ConfigEntryType, MediaType
17from torchaudio.transforms import SpectralCentroid
18
19from music_assistant.constants import VERBOSE_LOG_LEVEL
20from music_assistant.helpers.datetime import utc
21from music_assistant.helpers.util import is_arm, system_meets_requirements
22from music_assistant.models.audio_analysis import AudioAnalysisData, AudioAnalysisError
23from music_assistant.models.audio_analysis_provider import (
24 ACCUMULATING_ANALYSIS_MAX_DURATION_SECONDS,
25 AudioAnalysisProvider,
26)
27
28from .dbn_postprocessor import DBNDownBeatTracker
29from .feature_extractor import AdvancedBeatFeatureExtractor
30from .helpers import (
31 aggregate_series_to_bins,
32 calculate_overall_bpm,
33 compute_band_rms_frames,
34 decode_pcm_chunk_to_mono,
35)
36from .resources.skey_model import KEY_MAP as SKEY_KEY_MAP
37from .resources.skey_model import VQT, ChromaNet, CropCQT, load_skey_components
38from .vocal_activity import (
39 FIRERED_SAMPLE_RATE,
40 FireRedFbank,
41 infer_firered_chunk,
42 load_firered_components,
43 split_firered_features,
44 vocal_activity_probabilities,
45)
46
47if TYPE_CHECKING:
48 import numpy.typing as npt
49 from music_assistant_models.config_entries import ProviderConfig
50 from music_assistant_models.enums import ProviderFeature
51 from music_assistant_models.media_items import AudioFormat
52 from music_assistant_models.provider import ProviderManifest
53 from music_assistant_models.streamdetails import StreamDetails
54
55 from music_assistant.mass import MusicAssistant
56
57ANALYSIS_SAMPLE_RATE = 22050
58# Below the recommended thresholds the provider still runs, but we surface an
59# informational notice (see get_config_entries) as it may be tight under load.
60RECOMMENDED_RAM_GB = 6.0
61RECOMMENDED_CPU_CORES = 4
62# Beat This predicts a long track as fixed windows. These are the values the model was trained
63# and released with (30s at 50 fps, plus the loss-border frames its predictions are unreliable
64# on), so a windowed prediction is identical to a whole-track one. Do not tune them: a window
65# of another length puts the model off its training distribution across the whole window.
66BEAT_WINDOW_FRAMES = 1500
67BEAT_WINDOW_BORDER_FRAMES = 6
68BEAT_WINDOW_OVERLAP_MODE = "keep_first"
69# While a player streams, wait this many times a window's own compute time before starting the
70# next one, so beat inference does not occupy a core continuously.
71BEAT_WINDOW_PACE_RATIO = 1.0
72# Model failures such as one-off torch/hardware errors are often transient. Record them with
73# this retry horizon instead of a permanent row that blocks the track forever.
74MODEL_FAILURE_RETRY_DELAY = timedelta(hours=24)
75
76
77@dataclass
78class SmartFadesData:
79 """Per-session data for smart fades analysis."""
80
81 item_id: str
82 provider: str
83 input_audio_format: AudioFormat
84 block_samples: int
85 features: AdvancedBeatFeatureExtractor
86 resampler: soxr.ResampleStream | None = None
87 pcm_buffer: list[np.ndarray] = field(default_factory=list)
88 pcm_samples: int = 0
89 total_pcm_samples: int = 0
90 beats_feature_blocks: list[np.ndarray] = field(default_factory=list)
91 energy_chunks: list[np.ndarray] = field(default_factory=list)
92 centroid_chunks: list[np.ndarray] = field(default_factory=list)
93 frequency_band_chunks: dict[str, list[np.ndarray]] = field(default_factory=dict)
94 musical_key_feature_blocks: list[torch.Tensor] = field(default_factory=list)
95 vocal_resampler: soxr.ResampleStream | None = None
96 vocal_fbank: FireRedFbank | None = None
97 vocal_feature_blocks: list[np.ndarray] = field(default_factory=list)
98
99
100@dataclass(frozen=True, slots=True)
101class LoadedModels:
102 """The model components Smart Fades infers with; loaded and released as one set."""
103
104 beat_this: Spect2Frames
105 beat_this_post_processor: DBNDownBeatTracker
106 skey_vqt: VQT
107 skey_chromanet: ChromaNet
108 skey_crop: CropCQT
109 spectral_centroid: SpectralCentroid
110 firered: torch.nn.Module
111 firered_cmvn_means: npt.NDArray[np.float64]
112 firered_cmvn_inverse_std: npt.NDArray[np.float64]
113
114
115class SmartFadesProvider(AudioAnalysisProvider):
116 """Smart fades audio analysis provider using Beat This for beat tracking."""
117
118 max_analysis_duration = ACCUMULATING_ANALYSIS_MAX_DURATION_SECONDS
119 # v3: FireRed AED vocal activity
120 analysis_version = 3
121 has_unloadable_models = True
122
123 def __init__(
124 self,
125 mass: MusicAssistant,
126 manifest: ProviderManifest,
127 config: ProviderConfig,
128 supported_features: set[ProviderFeature],
129 ) -> None:
130 """Initialize the provider."""
131 super().__init__(mass, manifest, config, supported_features)
132 self._data: dict[str, SmartFadesData] = {}
133 self._device = "cpu"
134 # Populated by _load_models and cleared again by _free_models, so every use goes
135 # through _require_models rather than assuming the models are resident.
136 self._models: LoadedModels | None = None
137
138 async def get_config_entries(self) -> tuple[ConfigEntry, ...]:
139 """Return config entries for this provider."""
140 return (
141 ConfigEntry(
142 key="resource_warning",
143 type=ConfigEntryType.ALERT,
144 required=False,
145 hidden=system_meets_requirements(
146 min_memory_gb=RECOMMENDED_RAM_GB,
147 min_cpu_cores=RECOMMENDED_CPU_CORES,
148 ),
149 ),
150 )
151
152 async def handle_async_init(self) -> None:
153 """Handle async initialization of the provider; idle models are reloaded on demand."""
154 # Configure the inference runtime before loading any model (see the controller method).
155 self.mass.streams.audio_analysis.ensure_inference_runtime_configured()
156 await self._load_models()
157 self._models_loaded = True
158
159 async def process_pcm_chunk(
160 self,
161 session_id: str,
162 pcm_chunk: bytes,
163 ) -> None:
164 """Process a PCM chunk for beat tracking."""
165 data = self._data.get(session_id)
166 if not data:
167 return
168
169 pcm_mono = await self._run_offloaded(
170 decode_pcm_chunk_to_mono, data.input_audio_format, pcm_chunk
171 )
172 if pcm_mono.size == 0:
173 return
174
175 # Per-chunk VQT for key detection (skip short tail chunks)
176 if len(pcm_mono) >= data.input_audio_format.sample_rate:
177 await self._run_offloaded(
178 self._compute_musical_key_features,
179 pcm_mono,
180 data.input_audio_format.sample_rate,
181 data,
182 )
183
184 data.pcm_buffer.append(pcm_mono)
185 data.pcm_samples += len(pcm_mono)
186
187 # calculate features in 10s blocks to avoid cpu contention
188 if data.pcm_samples >= data.block_samples:
189 await self._process_block(data)
190
191 async def cancel(self, session_id: str) -> None:
192 """Cancel a beat tracking session."""
193 data = self._data.pop(session_id, None)
194 if data:
195 self._clear_session_data(data)
196 await super().cancel(session_id)
197
198 async def _load_models(self) -> None:
199 """Load the Beat This, S-KEY, and FireRed AED models into memory."""
200 self._models = await asyncio.to_thread(self._initialize_models)
201
202 def _free_models(self) -> None:
203 """Release the Beat This, S-KEY, and FireRed AED models."""
204 self._models = None
205
206 def _initialize_models(self) -> LoadedModels:
207 """Initialize ML models (runs in a thread to avoid blocking the event loop)."""
208 beat_this_model = Spect2Frames(checkpoint_path="small0", device=self._device)
209 # torch aarch64 wheels advertise fbgemm in supported_engines but its kernels are x86-only.
210 preference = ("qnnpack", "fbgemm") if is_arm() else ("fbgemm", "qnnpack")
211 supported_engines = torch.backends.quantized.supported_engines
212 quantized_engine = next((e for e in preference if e in supported_engines), None)
213 if quantized_engine is not None and torch.backends.quantized.engine != quantized_engine:
214 torch.backends.quantized.engine = quantized_engine
215 beat_this_model.model = torch.ao.quantization.quantize_dynamic( # type: ignore[no-untyped-call]
216 beat_this_model.model, {torch.nn.Linear}, dtype=torch.qint8
217 )
218 beat_this_post_processor = DBNDownBeatTracker(
219 beats_per_bar=[3, 4], min_bpm=55, max_bpm=215, fps=50
220 )
221 skey_vqt, skey_chromanet, skey_crop = load_skey_components(device=self._device)
222 spectral_centroid = SpectralCentroid(sample_rate=ANALYSIS_SAMPLE_RATE, hop_length=512)
223 firered_model, firered_cmvn_means, firered_cmvn_inverse_std = load_firered_components(
224 device=self._device
225 )
226 return LoadedModels(
227 beat_this=beat_this_model,
228 beat_this_post_processor=beat_this_post_processor,
229 skey_vqt=skey_vqt,
230 skey_chromanet=skey_chromanet,
231 skey_crop=skey_crop,
232 spectral_centroid=spectral_centroid,
233 firered=firered_model,
234 firered_cmvn_means=firered_cmvn_means,
235 firered_cmvn_inverse_std=firered_cmvn_inverse_std,
236 )
237
238 def _require_models(self) -> LoadedModels:
239 """Return the resident model set, failing the analysis when it has been unloaded."""
240 if self._models is None:
241 # The recorder that handles AudioAnalysisError does not log, so warn here or the
242 # unload race leaves nothing behind but a row in the failures table.
243 self.logger.warning("Models are not loaded; analysis will be retried later")
244 raise AudioAnalysisError(
245 "models are not loaded",
246 retry_at=utc() + MODEL_FAILURE_RETRY_DELAY,
247 )
248 return self._models
249
250 async def _start_analysis(
251 self,
252 session_id: str,
253 streamdetails: StreamDetails,
254 audio_format: AudioFormat,
255 ) -> bool:
256 """Start beat tracking analysis for a new track."""
257 if streamdetails.media_type != MediaType.TRACK:
258 # We only want to analyze tracks
259 return False
260
261 models = self._require_models()
262 block_seconds = 10.0
263
264 needs_resample = audio_format.sample_rate != ANALYSIS_SAMPLE_RATE
265 self._data[session_id] = SmartFadesData(
266 item_id=streamdetails.item_id,
267 provider=streamdetails.provider,
268 input_audio_format=audio_format,
269 block_samples=int(block_seconds * audio_format.sample_rate),
270 features=AdvancedBeatFeatureExtractor(
271 sample_rate=ANALYSIS_SAMPLE_RATE,
272 device=self._device,
273 offload=self._run_offloaded,
274 ),
275 resampler=soxr.ResampleStream(
276 in_rate=audio_format.sample_rate,
277 out_rate=ANALYSIS_SAMPLE_RATE,
278 num_channels=1,
279 dtype="float32",
280 )
281 if needs_resample
282 else None,
283 vocal_resampler=soxr.ResampleStream(
284 in_rate=audio_format.sample_rate,
285 out_rate=FIRERED_SAMPLE_RATE,
286 num_channels=1,
287 dtype="float32",
288 )
289 if audio_format.sample_rate != FIRERED_SAMPLE_RATE
290 else None,
291 vocal_fbank=FireRedFbank(
292 models.firered_cmvn_means,
293 models.firered_cmvn_inverse_std,
294 ),
295 )
296 self.logger.debug("Started beat tracking session %s", session_id)
297 return True
298
299 async def _finalize(self, session_id: str) -> AudioAnalysisData | None:
300 """Finalize beat tracking and store results."""
301 data = self._data.pop(session_id, None)
302 if not data:
303 return None
304
305 try:
306 if data.pcm_samples:
307 await self._process_block(data, last=True)
308 else:
309 # The vocal resampler and fbank still need an explicit end-of-input flush.
310 await self._run_offloaded(
311 self._compute_vocal_features,
312 np.empty(0, dtype=np.float32),
313 data,
314 True,
315 )
316
317 final_feats = await data.features.finalize()
318 if final_feats.size:
319 data.beats_feature_blocks.append(final_feats)
320 if not data.beats_feature_blocks:
321 return None
322
323 feats = np.concatenate(data.beats_feature_blocks, axis=0)
324 data.beats_feature_blocks.clear()
325 duration = data.total_pcm_samples / ANALYSIS_SAMPLE_RATE
326
327 all_vqt = None
328 if data.musical_key_feature_blocks:
329 all_vqt = torch.cat(data.musical_key_feature_blocks, dim=-1) # (1, 1, 84, T_total)
330 data.musical_key_feature_blocks.clear()
331
332 if data.vocal_feature_blocks:
333 vocal_features = np.concatenate(data.vocal_feature_blocks)
334 data.vocal_feature_blocks.clear()
335 else:
336 vocal_features = np.empty((0, 80), dtype=np.float32)
337
338 beat_key_result, vocal_activity = await self._run_final_inference(
339 feats,
340 all_vqt,
341 vocal_features,
342 duration,
343 )
344 beats, downbeats, beats_per_bar, key, mode = beat_key_result
345 return self._build_analysis(
346 data,
347 duration,
348 beats,
349 downbeats,
350 beats_per_bar,
351 key,
352 mode,
353 vocal_activity,
354 )
355 finally:
356 self._clear_session_data(data)
357
358 def _build_analysis(
359 self,
360 data: SmartFadesData,
361 duration: float,
362 beats: np.ndarray,
363 downbeats: np.ndarray,
364 beats_per_bar: int,
365 key: str | None,
366 mode: str | None,
367 vocal_activity: np.ndarray,
368 ) -> AudioAnalysisData:
369 """Build the final Smart Fades analysis payload."""
370 bpm = calculate_overall_bpm(beats)
371
372 # mean power per bin: point sampling aliases beat-rate ripple into the bins
373 rms_energy = None
374 energy_peak = 0.0
375 if data.energy_chunks:
376 energy_all = np.concatenate(data.energy_chunks)
377 if len(energy_all) >= 2:
378 rms_energy = aggregate_series_to_bins(energy_all, 1800, power=True)
379 energy_peak = float(rms_energy.max())
380 if energy_peak > 0:
381 rms_energy = rms_energy / energy_peak
382
383 spectral_centroid = None
384 if data.centroid_chunks:
385 centroid_all = np.concatenate(data.centroid_chunks)
386 if len(centroid_all) >= 2:
387 spectral_centroid = aggregate_series_to_bins(centroid_all, 1800)
388 # Zero out centroid where energy is negligible (noise dominates)
389 if rms_energy is not None:
390 spectral_centroid[rms_energy < 0.01] = 0.0
391
392 vocal_activity_bins = (
393 aggregate_series_to_bins(vocal_activity, 1800)
394 if vocal_activity.size
395 else np.zeros(1800, dtype=np.float32)
396 )
397 extra_data: dict[str, Any] = {"vocal_activity": vocal_activity_bins.tolist()}
398 if energy_peak > 0 and data.frequency_band_chunks:
399 extra_data["band_rms"] = {
400 name: (
401 aggregate_series_to_bins(np.concatenate(chunks), 1800, power=True) / energy_peak
402 ).tolist()
403 for name, chunks in data.frequency_band_chunks.items()
404 }
405
406 analysis = AudioAnalysisData(
407 bpm=bpm,
408 # the model stores plain float lists (numpy-free); convert the analysis arrays
409 beats=beats.tolist(),
410 downbeats=downbeats.tolist(),
411 duration=duration,
412 rms_energy=rms_energy.tolist() if rms_energy is not None else None,
413 spectral_centroid=(
414 spectral_centroid.tolist() if spectral_centroid is not None else None
415 ),
416 key=key,
417 mode=mode,
418 extra_data=extra_data,
419 beats_per_bar=beats_per_bar or None,
420 )
421 self.logger.debug(
422 "Beat analysis for %s: BPM=%.1f, %d beats, %d downbeats, key=%s",
423 data.item_id,
424 bpm,
425 len(beats),
426 len(downbeats),
427 f"{key} {mode}" if key else "unknown",
428 )
429 return analysis
430
431 async def _process_block(self, data: SmartFadesData, *, last: bool = False) -> None:
432 """Resample accumulated PCM buffer and extract features."""
433 start_time = time.perf_counter()
434 pcm_raw = (
435 np.concatenate(data.pcm_buffer) if data.pcm_buffer else np.empty(0, dtype=np.float32)
436 )
437 data.pcm_buffer.clear()
438 data.pcm_samples = 0
439
440 if data.resampler is not None:
441 pcm_22k = await self._run_offloaded(data.resampler.resample_chunk, pcm_raw, last)
442 else:
443 pcm_22k = pcm_raw
444
445 data.total_pcm_samples += len(pcm_22k)
446
447 if pcm_22k.size:
448 feats, _, _ = await asyncio.gather(
449 data.features.process_pcm(pcm_22k),
450 self._run_offloaded(
451 self._compute_energy_and_spectral_centroids,
452 pcm_22k,
453 data,
454 ),
455 self._run_offloaded(self._compute_vocal_features, pcm_raw, data, last),
456 )
457 else:
458 await self._run_offloaded(self._compute_vocal_features, pcm_raw, data, last)
459 feats = np.empty((0, 128), dtype=np.float32)
460
461 if feats.size:
462 data.beats_feature_blocks.append(feats)
463
464 elapsed_ms = (time.perf_counter() - start_time) * 1000
465 self.logger.log(VERBOSE_LOG_LEVEL, "Processed 10s of PCM chunks in %.1fms", elapsed_ms)
466
467 def _compute_energy_and_spectral_centroids(
468 self, pcm_22k: np.ndarray, data: SmartFadesData
469 ) -> None:
470 """Compute fine-resolution RMS energy and spectral centroid for a block."""
471 sr = ANALYSIS_SAMPLE_RATE
472 # RMS energy in 100ms windows, including partial final window
473 window_samples = sr // 10 # 2205 samples = 100ms
474 if len(pcm_22k) > 0:
475 n_full = len(pcm_22k) // window_samples
476 rms_list = []
477 if n_full > 0:
478 frames = pcm_22k[: n_full * window_samples].reshape(n_full, window_samples)
479 rms_list.append(np.sqrt(np.mean(frames**2, axis=1)))
480 remainder = len(pcm_22k) - n_full * window_samples
481 if remainder > 0:
482 tail = pcm_22k[n_full * window_samples :]
483 rms_list.append(np.array([np.sqrt(np.mean(tail**2))]))
484 if rms_list:
485 data.energy_chunks.append(np.concatenate(rms_list).astype(np.float32))
486
487 band_frames = compute_band_rms_frames(pcm_22k, sr, window_samples)
488 for name, frames in band_frames.items():
489 data.frequency_band_chunks.setdefault(name, []).append(frames)
490
491 # Spectral centroid: keep per-frame (hop_length=512, ~43 frames/s)
492 # Skip short tail buffers: STFT reflect-pad requires len > n_fft // 2.
493 spectral_centroid = self._require_models().spectral_centroid
494 if len(pcm_22k) >= spectral_centroid.n_fft:
495 pcm_tensor = torch.from_numpy(pcm_22k)
496 centroid_frames = spectral_centroid(pcm_tensor.unsqueeze(0)).squeeze(0).numpy()
497 # digitally-silent frames divide 0/0 into NaN; treat them as 0 Hz like
498 # other negligible-energy frames so no non-finite value is ever stored
499 np.nan_to_num(centroid_frames, copy=False, nan=0.0, posinf=0.0, neginf=0.0)
500 if len(centroid_frames) > 0:
501 data.centroid_chunks.append(centroid_frames.astype(np.float32))
502
503 def _compute_vocal_features(
504 self,
505 pcm_raw: np.ndarray,
506 data: SmartFadesData,
507 last: bool,
508 ) -> None:
509 """Resample source PCM to 16 kHz and extract FireRed fbank features."""
510 fbank = data.vocal_fbank
511 resampler = data.vocal_resampler
512 if fbank is None:
513 return
514 pcm_16k = resampler.resample_chunk(pcm_raw, last) if resampler is not None else pcm_raw
515 features = fbank.process(pcm_16k)
516 if features.size:
517 data.vocal_feature_blocks.append(features)
518 if last:
519 final_features = fbank.finalize()
520 if final_features.size:
521 data.vocal_feature_blocks.append(final_features)
522
523 async def _run_final_inference(
524 self,
525 beat_features: np.ndarray,
526 key_features: torch.Tensor | None,
527 vocal_features: np.ndarray,
528 duration: float,
529 ) -> tuple[tuple[np.ndarray, np.ndarray, int, str | None, str | None], np.ndarray]:
530 """Run beat/key and vocal inference branches; the first failure cancels the other."""
531 try:
532 async with asyncio.TaskGroup() as task_group:
533 beat_key_task = task_group.create_task(
534 self._infer_beats_and_key(beat_features, key_features)
535 )
536 vocal_task = task_group.create_task(
537 self._infer_vocal_activity(vocal_features, duration)
538 )
539 except ExceptionGroup as group:
540 # Unwrap for the base class; prefer the beat/key error since it decides
541 # permanent vs retryable failure recording.
542 beat_key_error = None if beat_key_task.cancelled() else beat_key_task.exception()
543 primary = beat_key_error or group.exceptions[0]
544 for error in group.exceptions:
545 if error is not primary:
546 self.logger.debug("FireRed vocal inference also failed: %s", error)
547 raise primary from primary.__cause__
548 return beat_key_task.result(), vocal_task.result()
549
550 async def _infer_beats_and_key(
551 self,
552 beat_features: np.ndarray,
553 key_features: torch.Tensor | None,
554 ) -> tuple[np.ndarray, np.ndarray, int, str | None, str | None]:
555 """Run beat inference followed by musical key inference."""
556 # Resolved before the beat stage: a reload or shutdown can free the models mid-run.
557 models = self._require_models()
558 beats, downbeats, beats_per_bar = await self._infer_beat_timings(beat_features)
559 if len(beats) < 2:
560 raise AudioAnalysisError("no rhythmic beat detected")
561 key, mode = await self._run_offloaded(
562 self._infer_musical_key, models.skey_chromanet, key_features
563 )
564 return beats, downbeats, beats_per_bar, key, mode
565
566 async def _infer_vocal_activity(
567 self,
568 features: np.ndarray,
569 duration: float,
570 ) -> np.ndarray:
571 """Run FireRed AED inference and return the 100 ms vocal timeline."""
572 model = self._require_models().firered
573 try:
574 compute_seconds = 0.0
575 chunks = []
576 for chunk, core_offset, core_length in split_firered_features(features):
577 chunk_probabilities, elapsed = await self._run_offloaded_timed(
578 infer_firered_chunk,
579 model,
580 chunk,
581 self._device,
582 )
583 compute_seconds += elapsed
584 chunks.append(chunk_probabilities[core_offset : core_offset + core_length])
585 frame_probabilities = (
586 np.concatenate(chunks) if chunks else np.empty((0, 3), dtype=np.float32)
587 )
588 probabilities = vocal_activity_probabilities(frame_probabilities, duration)
589 except asyncio.CancelledError:
590 raise
591 except Exception as err:
592 # Avoid permanently suppressing beat and key results for transient model failures.
593 raise AudioAnalysisError(
594 f"FireRed vocal inference failed: {err}",
595 retry_at=utc() + MODEL_FAILURE_RETRY_DELAY,
596 ) from err
597 self.logger.log(
598 VERBOSE_LOG_LEVEL,
599 "FireRed vocal inference: %.1fms compute over %d frames",
600 compute_seconds * 1000,
601 len(features),
602 )
603 return probabilities
604
605 def _compute_musical_key_features(
606 self, pcm_mono: np.ndarray, sample_rate: int, data: SmartFadesData
607 ) -> None:
608 """Extract VQT features for S-KEY key detection."""
609 models = self._require_models()
610 if sample_rate != ANALYSIS_SAMPLE_RATE:
611 pcm_mono = soxr.resample(pcm_mono, sample_rate, ANALYSIS_SAMPLE_RATE)
612 pcm_tensor = torch.from_numpy(pcm_mono)
613 with torch.inference_mode():
614 vqt_input = pcm_tensor.unsqueeze(0).unsqueeze(0) # (1, 1, samples)
615 vqt_out = models.skey_vqt(vqt_input) # (1, 1, n_bins, T)
616 cropped = models.skey_crop(vqt_out, torch.zeros(1)) # (1, 1, 84, T)
617 data.musical_key_feature_blocks.append(cropped.cpu())
618
619 def _infer_musical_key(
620 self, chromanet: torch.nn.Module, vqt_features: torch.Tensor | None
621 ) -> tuple[str | None, str | None]:
622 """
623 Run S-KEY ChromaNet inference to detect musical key.
624
625 :param chromanet: The ChromaNet module to run.
626 :param vqt_features: Accumulated VQT features, or None when the track had too few.
627 """
628 if vqt_features is None or vqt_features.shape[-1] < 128:
629 return None, None
630 start = time.perf_counter()
631 with torch.no_grad():
632 logits = chromanet(vqt_features.to(self._device))
633 key_idx = int(logits.argmax(dim=-1).item())
634 key_name = SKEY_KEY_MAP[key_idx] # e.g. "C# Major"
635 parts = key_name.split()
636 self.logger.log(
637 VERBOSE_LOG_LEVEL,
638 "ChromaNet key inference: %.1fms, detected key=%s %s",
639 (time.perf_counter() - start) * 1000,
640 parts[0],
641 parts[1],
642 )
643 return parts[0], parts[1].lower()
644
645 async def _infer_beat_timings(self, feats: np.ndarray) -> tuple[np.ndarray, np.ndarray, int]:
646 """
647 Run Beat This model inference to detect beat/downbeat timings and the meter.
648
649 :param feats: Log-mel features for the whole track, shaped (frames, mel bins).
650 """
651 # Resolved once and passed down: a reload or shutdown can free the models while the
652 # windows below are still being dispatched.
653 models = self._require_models()
654
655 spect = torch.from_numpy(feats).to(self._device)
656 windows, starts = split_piece(
657 spect,
658 BEAT_WINDOW_FRAMES,
659 border_size=BEAT_WINDOW_BORDER_FRAMES,
660 avoid_short_end=True,
661 )
662 predictions = []
663 model_seconds = 0.0
664 for window in windows:
665 prediction, elapsed = await self._run_offloaded_timed(
666 self._infer_beat_window, models.beat_this.model, window
667 )
668 predictions.append(prediction)
669 model_seconds += elapsed
670 await self._pace_beat_windows(elapsed)
671
672 (beats, downbeats, beats_per_bar), post_seconds = await self._run_offloaded_timed(
673 self._decode_beat_timings,
674 models.beat_this_post_processor,
675 predictions,
676 starts,
677 len(spect),
678 )
679 self.logger.log(
680 VERBOSE_LOG_LEVEL,
681 "Model inference: %.1fms compute over %d windows, postprocessing: %.1fms, "
682 "detected %d beats, %d downbeats",
683 model_seconds * 1000,
684 len(windows),
685 post_seconds * 1000,
686 len(beats),
687 len(downbeats),
688 )
689 return beats, downbeats, beats_per_bar
690
691 async def _pace_beat_windows(self, window_seconds: float) -> None:
692 """
693 Idle for as long as the beat inference window that just finished took to compute.
694
695 Only while a player streams; idle and background analysis run at full speed.
696
697 :param window_seconds: Compute time of the window that just finished.
698 """
699 if not self.mass.streams.audio_analysis.playback_active():
700 return
701 await asyncio.sleep(window_seconds * BEAT_WINDOW_PACE_RATIO)
702
703 @staticmethod
704 def _infer_beat_window(model: torch.nn.Module, window: torch.Tensor) -> dict[str, torch.Tensor]:
705 """
706 Run one Beat This window and return its beat and downbeat logits.
707
708 :param model: The Beat This module to run.
709 :param window: One window of log-mel features, shaped (frames, mel bins).
710 """
711 # inference_mode is thread-local, so it has to be entered on the worker thread.
712 with torch.inference_mode():
713 prediction = model(window.unsqueeze(0))
714 return {"beat": prediction["beat"][0], "downbeat": prediction["downbeat"][0]}
715
716 def _decode_beat_timings(
717 self,
718 post_processor: DBNDownBeatTracker,
719 predictions: list[dict[str, torch.Tensor]],
720 starts: np.ndarray,
721 total_frames: int,
722 ) -> tuple[np.ndarray, np.ndarray, int]:
723 """
724 Stitch per-window logits back together and decode them into beat timings.
725
726 :param post_processor: The DBN decoder to run on the stitched activations.
727 :param predictions: Per-window beat/downbeat logits, in window order.
728 :param starts: Frame offset of each window, as returned by split_piece.
729 :param total_frames: Frame count of the whole track.
730 """
731 with torch.inference_mode():
732 beat_logits, downbeat_logits = aggregate_prediction(
733 predictions,
734 starts,
735 total_frames,
736 BEAT_WINDOW_FRAMES,
737 BEAT_WINDOW_BORDER_FRAMES,
738 BEAT_WINDOW_OVERLAP_MODE,
739 self._device,
740 )
741 dbn_out, beats_per_bar = post_processor(
742 self._beat_activations(beat_logits.float(), downbeat_logits.float())
743 )
744 beats = dbn_out[:, 0]
745 downbeats = dbn_out[dbn_out[:, 1] == 1, 0]
746 return beats, downbeats, beats_per_bar
747
748 @staticmethod
749 def _beat_activations(beat_logits: torch.Tensor, downbeat_logits: torch.Tensor) -> np.ndarray:
750 """
751 Convert beat/downbeat logits into the (T, 2) activations the DBN expects.
752
753 :param beat_logits: Per-frame beat logits for the whole track.
754 :param downbeat_logits: Per-frame downbeat logits for the whole track.
755 """
756 beat_prob = torch.sigmoid(beat_logits).cpu().numpy()
757 downbeat_prob = torch.sigmoid(downbeat_logits).cpu().numpy()
758 epsilon = 1e-5
759 beat_prob = beat_prob * (1 - epsilon) + epsilon / 2
760 downbeat_prob = downbeat_prob * (1 - epsilon) + epsilon / 2
761 return np.column_stack(
762 [
763 np.maximum(beat_prob - downbeat_prob, epsilon / 2),
764 downbeat_prob,
765 ]
766 )
767
768 def _clear_session_data(self, data: SmartFadesData) -> None:
769 """Release all state retained for an analysis session."""
770 data.pcm_buffer.clear()
771 data.beats_feature_blocks.clear()
772 data.energy_chunks.clear()
773 data.centroid_chunks.clear()
774 data.frequency_band_chunks.clear()
775 data.musical_key_feature_blocks.clear()
776 data.vocal_feature_blocks.clear()
777 data.resampler = None
778 data.vocal_resampler = None
779 data.vocal_fbank = None
780