/
/
/
1"""Tests for FireRed AED vocal activity inference."""
2
3from __future__ import annotations
4
5import math
6from unittest.mock import patch
7
8import numpy as np
9import pytest
10import torch
11
12from music_assistant.providers.smart_fades.vocal_activity import (
13 FIRERED_INFERENCE_CONTEXT_FRAMES,
14 FIRERED_MAX_INFERENCE_FRAMES,
15 FIRERED_PARAMETER_COUNT,
16 FireRedFbank,
17 infer_firered_chunk,
18 load_firered_components,
19 quantize_pcm_for_firered,
20 split_firered_features,
21 vocal_activity_probabilities,
22)
23
24
25def _reference_pcm() -> np.ndarray:
26 """Return deterministic normalized PCM derived from signed 16-bit samples."""
27 time = np.arange(16000 * 3.2, dtype=np.float64) / 16000
28 pcm_int16 = np.rint(
29 12000 * np.sin(2 * np.pi * 220 * time) + 4000 * np.sin(2 * np.pi * 997 * time)
30 ).astype(np.int16)
31 result: np.ndarray = pcm_int16.astype(np.float32) / 32768.0
32 return result
33
34
35def test_model_loads_weights_only_with_expected_parameter_count() -> None:
36 """The bundled state dict loads safely into the expected FireRed AED model."""
37 with patch(
38 "music_assistant.providers.smart_fades.vocal_activity.torch.load",
39 wraps=torch.load,
40 ) as load:
41 model, means, inverse_std = load_firered_components()
42
43 assert load.call_args.kwargs["weights_only"] is True
44 assert sum(parameter.numel() for parameter in model.parameters()) == (FIRERED_PARAMETER_COUNT)
45 assert means.shape == inverse_std.shape == (80,)
46
47
48def test_pcm_quantization_matches_reference_int16_scale() -> None:
49 """Normalized PCM is rounded and clipped to the FireRed signed 16-bit scale."""
50 pcm = np.array(
51 [-1.1, -1.0, -0.5, -0.5 / 32768, 0.5 / 32768, 0.5, 1.0, 1.1],
52 dtype=np.float32,
53 )
54
55 quantized = quantize_pcm_for_firered(pcm)
56
57 np.testing.assert_array_equal(
58 quantized,
59 np.array(
60 [-32768, -32768, -16384, 0, 0, 16384, 32767, 32767],
61 dtype=np.float32,
62 ),
63 )
64
65
66def test_streaming_fbank_cmvn_matches_reference_across_chunks() -> None:
67 """Arbitrary chunk boundaries preserve exact online fbank and CMVN features."""
68 pcm = _reference_pcm()
69 _, means, inverse_std = load_firered_components()
70
71 one_shot_extractor = FireRedFbank(means, inverse_std)
72 one_shot = np.concatenate([one_shot_extractor.process(pcm), one_shot_extractor.finalize()])
73
74 streaming_extractor = FireRedFbank(means, inverse_std)
75 blocks = [
76 features
77 for chunk in np.array_split(pcm, 37)
78 if (features := streaming_extractor.process(chunk)).size
79 ]
80 blocks.append(streaming_extractor.finalize())
81 streaming = np.concatenate(blocks)
82
83 np.testing.assert_array_equal(streaming, one_shot)
84 assert one_shot.shape == (318, 80)
85 np.testing.assert_allclose(
86 one_shot[[0, 1, 50, -1]][:, [0, 10, 40, 79]],
87 [
88 [1.161562, 0.5089607, -2.330163, -1.5592412],
89 [0.8259135, 0.50769424, -2.2977846, -1.700644],
90 [1.1541193, 0.50818205, -2.3833172, -1.7334999],
91 [0.9781029, 0.5056236, -2.3772168, -1.7556578],
92 ],
93 rtol=1e-4,
94 atol=1e-6,
95 )
96
97
98def test_deterministic_probability_parity() -> None:
99 """Bundled inference matches deterministic FireRed reference probabilities."""
100 pcm = _reference_pcm()
101 model, means, inverse_std = load_firered_components()
102 extractor = FireRedFbank(means, inverse_std)
103 features = np.concatenate([extractor.process(pcm), extractor.finalize()])
104
105 probabilities = infer_firered_chunk(model, features)
106 vocal_activity = vocal_activity_probabilities(probabilities, 3.2)
107
108 np.testing.assert_allclose(
109 probabilities[[0, 1, 50, 150, -1]],
110 [
111 [8.8199326e-05, 6.3637628e-05, 1.9249251e-01],
112 [8.2634397e-05, 6.1347811e-05, 1.5615767e-01],
113 [5.5029866e-04, 5.3342991e-04, 9.9471587e-01],
114 [5.2045441e-05, 4.0918807e-05, 9.9911863e-01],
115 [5.1496401e-05, 1.5160468e-04, 9.8695940e-01],
116 ],
117 rtol=1e-4,
118 atol=1e-6,
119 )
120 np.testing.assert_allclose(
121 vocal_activity[:5],
122 [8.7163455e-05, 2.1780634e-04, 4.3790121e-04, 5.1874004e-04, 5.9412332e-04],
123 rtol=1e-4,
124 atol=1e-7,
125 )
126
127
128@pytest.mark.parametrize(
129 ("frame_count", "expected"),
130 [
131 (FIRERED_MAX_INFERENCE_FRAMES, [(30_000, 0, 30_000)]),
132 (FIRERED_MAX_INFERENCE_FRAMES + 1, [(30_001, 0, 30_000), (161, 160, 1)]),
133 ],
134)
135def test_long_input_split_boundaries(
136 frame_count: int,
137 expected: list[tuple[int, int, int]],
138) -> None:
139 """Long inference adds and discards the required context at the 30k boundary."""
140 features = np.zeros((frame_count, 80), dtype=np.float32)
141
142 chunks = [
143 (len(chunk), core_offset, core_length)
144 for chunk, core_offset, core_length in split_firered_features(features)
145 ]
146
147 assert FIRERED_INFERENCE_CONTEXT_FRAMES == 160
148 assert chunks == expected
149
150
151def test_long_input_split_matches_full_inference_at_boundary() -> None:
152 """Context-trimmed chunks preserve full-model output at the split boundary."""
153 model, _, _ = load_firered_components()
154 # Keep the split far enough from the end that context cannot span the full input.
155 features = np.random.default_rng(7).normal(size=(30_500, 80)).astype(np.float32)
156
157 full = infer_firered_chunk(model, features)
158 split = np.concatenate(
159 [
160 infer_firered_chunk(model, chunk)[core_offset : core_offset + core_length]
161 for chunk, core_offset, core_length in split_firered_features(features)
162 ]
163 )
164
165 # Different tensor lengths cause minor float32 non-associativity.
166 np.testing.assert_allclose(split, full, atol=2e-6)
167
168
169def test_vocal_timeline_uses_max_speech_singing_and_exact_length() -> None:
170 """Each 100 ms value averages max(speech, singing) and pads the snipped tail."""
171 frame_probabilities = np.zeros((18, 3), dtype=np.float32)
172 frame_probabilities[:10, 0] = np.linspace(0.0, 0.9, 10)
173 frame_probabilities[:10, 1] = np.linspace(0.9, 0.0, 10)
174 frame_probabilities[10:, 0] = 1.2
175 frame_probabilities[10:, 1] = -0.2
176
177 probabilities = vocal_activity_probabilities(frame_probabilities, 0.21)
178
179 expected_first = np.maximum(
180 frame_probabilities[:10, 0],
181 frame_probabilities[:10, 1],
182 ).mean()
183 assert len(probabilities) == math.ceil(0.21 / 0.1)
184 assert probabilities[0] == pytest.approx(expected_first)
185 assert probabilities[1:].tolist() == pytest.approx([1.0, 1.0])
186 assert np.isfinite(probabilities).all()
187 assert np.all((probabilities >= 0.0) & (probabilities <= 1.0))
188
189
190def test_vocal_timeline_rejects_non_finite_model_output() -> None:
191 """Non-finite FireRed output cannot reach persisted extra_data."""
192 frame_probabilities = np.zeros((10, 3), dtype=np.float32)
193 frame_probabilities[4, 1] = np.nan
194
195 with pytest.raises(ValueError, match="non-finite"):
196 vocal_activity_probabilities(frame_probabilities, 0.1)
197