/
/
/
1"""
2Audio tap for the MilkDrop visualizer: waveform frames straight from playback.
3
4Reads the decoded PCM that the streams controller already buffers for the item
5a player is playing, so any player produces a waveform whatever protocol it
6renders over. The read is passive - it takes no audio away from playback and
7changes nothing about grouping or output protocol - so watching the visualizer
8never interrupts what is playing.
9
10Frames carry a play-at timestamp in the relay's clock domain. The tap anchors
11the track's media timeline to that clock from the queue's reported position,
12so a viewer draws each frame around the time it is audible. How closely that
13matches depends on how precisely the player reports its position.
14"""
15
16from __future__ import annotations
17
18import asyncio
19import dataclasses
20import struct
21import time
22from collections import deque
23from typing import TYPE_CHECKING, cast
24
25import numpy as np
26from music_assistant_models.enums import PlaybackState
27from music_assistant_models.media_items import MediaItemPalette
28from orjson import dumps
29
30from music_assistant.controllers.streams.audio_analysis import SMART_FADES_ANALYSIS_DOMAIN
31from music_assistant.controllers.streams.audio_buffer import AudioBufferDiscarded, AudioBufferEOF
32
33if TYPE_CHECKING:
34 from music_assistant_models.media_items import AudioFormat
35 from music_assistant_models.player_queue import PlayerQueue
36 from music_assistant_models.queue_item import QueueItem
37
38 from music_assistant.controllers.streams.audio_buffer import AudioBuffer
39 from music_assistant.models.audio_analysis import AudioAnalysisData
40 from music_assistant.models.player import Player
41
42 from .provider import MilkdropVisualizerProvider
43
44WAVE_SAMPLES = 1024
45CONF_COLOR_TINT = "color_tint"
46DEFAULT_COLOR_TINT = True
47# Derived from the model, so a field added upstream is forwarded automatically.
48COLOR_FIELDS = tuple(field.name for field in dataclasses.fields(MediaItemPalette))
49# How far ahead of the audible playhead the tap reads. Viewers schedule frames
50# by timestamp, so a lead is what lets them draw on time; it costs nothing,
51# since this audio is buffered already.
52LEAD_SECONDS = 5.0
53# Gap between where the anchor says the playhead is and where the queue reports
54# it that means the audio moved (a seek) rather than the player simply reporting
55# its position coarsely. Well above the ~1s quantization of the coarsest
56# reporters (DLNA), whose jitter would otherwise restart the frame flow.
57# Scaled by playback speed where compared: report jitter lives in wall-clock
58# time and inflates by the speed factor on its way into media time.
59RESYNC_THRESHOLD_SECONDS = 3.0
60# Poll interval while there is nothing to read: idle player, or the tap having
61# read as far ahead as it may.
62IDLE_POLL_SECONDS = 0.5
63# The neural beat tracker lands ~5-10s into a track, so a track that has beats
64# at all rarely has them at its first frame. Capped so a track that will never
65# have them stops asking.
66BEAT_RETRY_SECONDS = 3.0
67BEAT_RETRY_ATTEMPTS = 30
68# Frames a tap keeps to replay to a viewer that attaches mid-track, and the
69# ceiling on one viewer's outbound queue. The ring must span far more than
70# LEAD_SECONDS: on a track longer than the buffer's retained window, eviction
71# follows the player's stream pull (readrate 2x for HTTP players, ~30s commit
72# lead for Sendspin), so the tap is forced to read - and stamp - audio well
73# ahead of the audible playhead. The ring bridges that gap for attaching
74# viewers: ~95s at ~43 frames/s of ~1KB each (~4MB per tap). Beyond it the
75# audio is already evicted server-side, so no ring size can help; long tracks
76# spend their pinned phase there and that is an accepted limitation.
77RING_FRAMES = 4096
78VIEWER_QUEUE_FRAMES = 1024
79
80# Wire tags, matching the format documented in relay.py.
81WAVE_FRAME_TAG = 22
82BEAT_FRAME_TAG = 17
83
84
85def server_now_us() -> int:
86 """Return the relay's clock in microseconds, the domain frame timestamps live in."""
87 # Monotonic: a viewer only ever needs the server's clock to be consistent
88 # with itself, and a wall clock stepping under NTP would strand every frame
89 # already scheduled.
90 return int(time.monotonic() * 1_000_000)
91
92
93def pack_wave_frame(timestamp_us: int, samples: bytes) -> bytes:
94 """Pack one waveform tail for the wire."""
95 return struct.pack(">Bq", WAVE_FRAME_TAG, timestamp_us) + samples
96
97
98def pack_beat_frame(timestamp_us: int, is_downbeat: bool) -> bytes:
99 """Pack one beat schedule entry for the wire."""
100 return struct.pack(">BqB", BEAT_FRAME_TAG, timestamp_us, 1 if is_downbeat else 0)
101
102
103def pcm_to_mono(data: bytes, pcm_format: AudioFormat) -> np.ndarray:
104 """
105 Return a PCM chunk as mono float32 in -1.0..1.0.
106
107 :param data: Raw interleaved PCM as the playback buffer holds it.
108 :param pcm_format: The buffer's PCM format, giving bit depth and channel count.
109 """
110 bit_depth = pcm_format.bit_depth
111 if bit_depth == 24:
112 # No numpy dtype covers packed 24-bit, so assemble the samples by byte.
113 packed = np.frombuffer(data, dtype=np.uint8)
114 packed = packed[: packed.size - packed.size % 3].reshape(-1, 3).astype(np.int32)
115 raw = packed[:, 0] | packed[:, 1] << 8 | packed[:, 2] << 16
116 raw[raw >= 1 << 23] -= 1 << 24
117 scale = float(1 << 23)
118 else:
119 dtype = "<i2" if bit_depth == 16 else "<i4"
120 width = np.dtype(dtype).itemsize
121 raw = np.frombuffer(data[: len(data) - len(data) % width], dtype=dtype)
122 scale = float(1 << (width * 8 - 1))
123 channels = max(1, pcm_format.channels)
124 mono: np.ndarray
125 if channels > 1:
126 # Fold to mono in one pass, straight to float32: converting the whole
127 # interleaved chunk first would cost twice the memory for no gain.
128 # A truncated chunk drops its dangling frame rather than failing here.
129 raw = raw[: raw.size - raw.size % channels]
130 mono = raw.reshape(-1, channels).mean(axis=1, dtype=np.float32)
131 else:
132 mono = raw.astype(np.float32)
133 mono /= scale
134 return mono
135
136
137def palette_payload(palette: MediaItemPalette | None) -> dict[str, list[int] | None]:
138 """
139 Return a color@v1 payload for a track palette.
140
141 A track without a palette yields every field as null, so a viewer drops the
142 previous track's tint rather than keeping it over the new one.
143
144 :param palette: The palette resolved for the artwork now showing, if any.
145 """
146 payload: dict[str, list[int] | None] = {}
147 for name in COLOR_FIELDS:
148 value = getattr(palette, name, None) if palette is not None else None
149 payload[name] = list(value) if value else None
150 return payload
151
152
153class ViewerQueue:
154 """
155 Outbound queue for one viewer.
156
157 Bounded so a stalled browser cannot stall the tap, but control messages
158 (stream/clear, stream/end) are never dropped: losing one would leave the
159 viewer animating stale audio after a seek or track change.
160 """
161
162 def __init__(self, capacity: int = VIEWER_QUEUE_FRAMES) -> None:
163 """
164 Initialize the queue.
165
166 :param capacity: Maximum number of pending items before eviction kicks in.
167 """
168 self._items: deque[bytes | str] = deque()
169 self._capacity = capacity
170 self._wakeup = asyncio.Event()
171
172 def push(self, item: bytes | str) -> None:
173 """Enqueue an item, evicting the oldest waveform frame when full."""
174 if len(self._items) >= self._capacity:
175 for index, queued in enumerate(self._items):
176 if isinstance(queued, bytes):
177 del self._items[index]
178 break
179 else:
180 self._items.popleft()
181 self._items.append(item)
182 self._wakeup.set()
183
184 async def get(self) -> bytes | str:
185 """Wait for and return the next item."""
186 while not self._items:
187 self._wakeup.clear()
188 await self._wakeup.wait()
189 return self._items.popleft()
190
191
192@dataclasses.dataclass
193class TrackCursor:
194 """Where a tap has read to in the current track, and how its media time maps to the clock."""
195
196 item_id: str
197 # Clock time at which this track's media time zero was (or will be) audible.
198 anchor_us: int
199 # Next 1-second buffer chunk to read; chunk N is second N of the track.
200 next_chunk: int
201 # Samples left over from the previous chunk, and the media time they start at.
202 carry: np.ndarray
203 carry_media: float
204 # Media seconds per wall-clock second (atempo, audiobooks/podcasts).
205 speed: float = 1.0
206
207 def playhead(self) -> float:
208 """Return where the anchor says the audible playhead is now, in media seconds."""
209 return (server_now_us() - self.anchor_us) / 1_000_000 * self.speed
210
211 def media_to_clock_us(self, media_seconds: float) -> int:
212 """Return the clock time at which a media position becomes audible."""
213 return self.anchor_us + int(media_seconds / self.speed * 1_000_000)
214
215
216class Tap:
217 """One reader of a player's audio, shared by every viewer watching that player."""
218
219 def __init__(self, player_id: str) -> None:
220 """
221 Initialize the tap.
222
223 :param player_id: The player whose audio this tap follows.
224 """
225 self.player_id = player_id
226 self.queues: set[ViewerQueue] = set()
227 # Beat frames with their scheduled timestamps, so viewers that attach
228 # mid-track still receive the rest of the track's downbeats.
229 self.beats: deque[tuple[int, bytes]] = deque(maxlen=4096)
230 # Recent packed waveform frames, replayed to a connecting viewer so it
231 # has something to draw before the tap reaches its next chunk.
232 self.ring: deque[bytes] = deque(maxlen=RING_FRAMES)
233 # Latest color@v1 fields, replayed to viewers that attach mid-track.
234 self.last_color: dict[str, list[int] | None] = {}
235 # Beat analysis already fetched for the current item, so a re-anchor
236 # (a seek) rebuilds the schedule without querying again. Positive only:
237 # a cached miss would suppress analysis that lands late in the track.
238 self.beats_analysis: tuple[str, AudioAnalysisData] | None = None
239 # Set by the relay when a viewer attaches and finds only future-stamped
240 # frames; the reader consumes it by re-anchoring at the playhead.
241 self.realign_requested = False
242 self.task: asyncio.Task[None] | None = None
243 self.beats_task: asyncio.Task[None] | None = None
244
245 def fan_out(self, frame: bytes | str) -> None:
246 """Deliver a packed frame to every attached viewer queue."""
247 for queue in self.queues:
248 queue.push(frame)
249
250 def apply_color(self, payload: dict[str, list[int] | None]) -> None:
251 """Cache a color@v1 payload and fan it out."""
252 self.last_color = payload
253 self.fan_out(dumps({"type": "color", "payload": payload}).decode())
254
255 def has_only_future_frames(self) -> bool:
256 """
257 Return whether every buffered waveform frame is stamped ahead of now.
258
259 True when production is pinned at the buffer's eviction edge, ahead of
260 the audible playhead: the ring then holds nothing a fresh viewer could
261 draw yet, and re-anchoring at the playhead serves it better than a
262 replay would.
263 """
264 if not self.ring:
265 return False
266 timestamp_us: int = struct.unpack_from(">q", self.ring[0], 1)[0]
267 return timestamp_us > server_now_us()
268
269 def reset(self, message: str) -> None:
270 """Drop everything scheduled from a timeline that no longer applies."""
271 # a hydration still in flight would land beats for that dead timeline
272 if self.beats_task is not None:
273 self.beats_task.cancel()
274 self.beats_task = None
275 self.ring.clear()
276 self.beats.clear()
277 self.fan_out(message)
278
279
280class TapManager:
281 """Creates, shares and tears down the audio taps."""
282
283 def __init__(self, provider: MilkdropVisualizerProvider) -> None:
284 """
285 Initialize the tap manager.
286
287 :param provider: The loaded MilkDrop visualizer provider instance.
288 """
289 self.mass = provider.mass
290 self.provider = provider
291 self.logger = provider.logger.getChild("tap")
292 # One shared tap per player id, refcounted by viewer queues.
293 self._taps: dict[str, Tap] = {}
294 self._lock = asyncio.Lock()
295
296 async def acquire(self, player: Player) -> Tap:
297 """
298 Return the shared tap for a player, creating it on first use.
299
300 :param player: The player whose audio to follow.
301 """
302 async with self._lock:
303 if (existing := self._taps.get(player.player_id)) is not None:
304 return existing
305 tap = Tap(player.player_id)
306 self._taps[player.player_id] = tap
307 tap.task = self.mass.create_task(self._run(tap))
308 self.logger.info("Waveform tap following %s", player.display_name)
309 return tap
310
311 def schedule_release(self, player_id: str) -> None:
312 """
313 Tear down a tap whose last viewer just left.
314
315 Keyed per player with abort_existing, so a player never accumulates
316 releases: an earlier one could otherwise tear down a tap that a later
317 viewer created.
318
319 :param player_id: The player whose tap may now be idle.
320 """
321 self.mass.create_task(
322 self._release(player_id),
323 task_id=f"milkdrop_release_{player_id}",
324 abort_existing=True,
325 )
326
327 async def close(self) -> None:
328 """Tear down every live tap."""
329 async with self._lock:
330 for tap in self._taps.values():
331 self._stop(tap)
332 self._taps.clear()
333
334 def pending_beat_frames(self, tap: Tap) -> list[bytes]:
335 """
336 Return the tap's beat frames that are still in the future.
337
338 :param tap: The tap whose beat schedule to filter.
339 """
340 now_us = server_now_us()
341 return [frame for timestamp_us, frame in tap.beats if timestamp_us > now_us]
342
343 async def _release(self, player_id: str) -> None:
344 """Stop and forget a tap as soon as its last viewer goes."""
345 async with self._lock:
346 tap = self._taps.get(player_id)
347 if tap is None or tap.queues:
348 return
349 self._taps.pop(player_id, None)
350 self._stop(tap)
351 self.logger.info("Waveform tap for %s removed (viewers gone)", player_id)
352
353 def _stop(self, tap: Tap) -> None:
354 """Cancel a tap's reader and anything still working for it."""
355 for task in (tap.task, tap.beats_task):
356 if task is not None:
357 task.cancel()
358 tap.task = None
359 tap.beats_task = None
360
361 async def _run(self, tap: Tap) -> None:
362 """Read the player's audio for as long as the tap lives, packing frames for its viewers."""
363 cursor: TrackCursor | None = None
364 while True:
365 try:
366 cursor = await self._read_once(tap, cursor)
367 except Exception as err:
368 # a source that failed mid-track is a playback problem, not a
369 # reason for this tap to stop following the player
370 self.logger.debug("Tap for %s could not read: %s", tap.player_id, err)
371 cursor = None
372 await asyncio.sleep(IDLE_POLL_SECONDS)
373
374 async def _read_once(self, tap: Tap, cursor: TrackCursor | None) -> TrackCursor | None:
375 """
376 Advance a tap by at most one buffer chunk.
377
378 :param tap: The tap being fed.
379 :param cursor: The cursor from the previous pass, if it still has one.
380 :return: The cursor to carry into the next pass, or None to start over.
381 """
382 source = self._playing_source(tap.player_id)
383 if source is None:
384 if cursor is not None:
385 tap.reset('{"type": "stream/end"}')
386 await asyncio.sleep(IDLE_POLL_SECONDS)
387 return None
388 queue, item, buffer = source
389 self._sync_color(tap)
390 if tap.realign_requested:
391 # a viewer found only future-stamped frames; drop the cursor so the
392 # re-anchor below restarts at the playhead, but only while that
393 # chunk is still retained (past the edge a realign helps nobody)
394 tap.realign_requested = False
395 if queue.corrected_elapsed_time >= buffer.first_buffered_chunk:
396 cursor = None
397 cursor = self._align(
398 tap, cursor, item, queue.corrected_elapsed_time, buffer, queue.playback_speed
399 )
400 # Stay ahead of the listener, but never behind the buffer's retained
401 # window: a rolling (radio) buffer discards as playback consumes it, and
402 # what it is about to drop is the last chance to read that audio.
403 if (
404 cursor.next_chunk > cursor.playhead() + LEAD_SECONDS
405 and cursor.next_chunk > buffer.first_buffered_chunk
406 ):
407 await asyncio.sleep(IDLE_POLL_SECONDS)
408 return cursor
409 try:
410 pcm = await buffer.read_chunk_for_analysis(cursor.next_chunk)
411 except AudioBufferEOF:
412 # read past the end of a track that is still finishing; the next
413 # item takes over as soon as the queue moves on
414 await asyncio.sleep(IDLE_POLL_SECONDS)
415 return cursor
416 except AudioBufferDiscarded:
417 # the retained window moved past us (a stalled tap, or a rolling
418 # buffer outrunning it); pick the timeline up again where it is now
419 await asyncio.sleep(IDLE_POLL_SECONDS)
420 return None
421 self._emit_chunk(tap, cursor, pcm, buffer.pcm_format)
422 return cursor
423
424 def _playing_source(self, player_id: str) -> tuple[PlayerQueue, QueueItem, AudioBuffer] | None:
425 """Return the queue, item and PCM buffer of what a player is playing right now."""
426 queue = self.mass.player_queues.get_active_queue(player_id)
427 if queue is None or queue.state != PlaybackState.PLAYING:
428 return None
429 item = queue.current_item
430 if item is None or item.streamdetails is None:
431 return None
432 # An external source (a provider streaming straight to the device) has
433 # no buffer here, so there is no audio for us to read.
434 buffer = cast("AudioBuffer | None", item.streamdetails.buffer)
435 if buffer is None:
436 return None
437 return queue, item, buffer
438
439 def _align(
440 self,
441 tap: Tap,
442 cursor: TrackCursor | None,
443 item: QueueItem,
444 playhead: float,
445 buffer: AudioBuffer,
446 speed: float = 1.0,
447 ) -> TrackCursor:
448 """
449 Return a cursor whose timeline still matches what the player is playing.
450
451 A new track, a seek, a resume or a speed change re-anchors the media
452 timeline to the relay clock; so does falling behind the buffer's
453 retained window.
454
455 :param tap: The tap being fed.
456 :param cursor: The cursor in use, if the tap already has one.
457 :param item: The queue item now playing.
458 :param playhead: Media position the queue reports for it, in seconds.
459 :param buffer: The item's PCM buffer.
460 :param speed: Playback speed the queue plays the item at.
461 """
462 oldest = buffer.first_buffered_chunk
463 if (
464 cursor is not None
465 and cursor.item_id == item.queue_item_id
466 and cursor.speed == speed
467 and cursor.next_chunk >= oldest
468 and abs(cursor.playhead() - playhead) <= RESYNC_THRESHOLD_SECONDS * speed
469 ):
470 return cursor
471 start_chunk = max(int(max(0.0, playhead)), oldest)
472 cursor = TrackCursor(
473 item_id=item.queue_item_id,
474 anchor_us=server_now_us() - int(playhead / speed * 1_000_000),
475 next_chunk=start_chunk,
476 carry=np.zeros(0, dtype=np.float32),
477 carry_media=float(start_chunk),
478 speed=speed,
479 )
480 tap.reset('{"type": "stream/clear"}')
481 self._schedule_beats(tap, item, cursor.anchor_us, speed)
482 return cursor
483
484 def _emit_chunk(
485 self, tap: Tap, cursor: TrackCursor, pcm: bytes, pcm_format: AudioFormat
486 ) -> None:
487 """Turn one second of PCM into packed waveform frames and fan them out."""
488 sample_rate = pcm_format.sample_rate
489 mono = pcm_to_mono(pcm, pcm_format)
490 chunk_media = float(cursor.next_chunk)
491 if abs(cursor.carry_media + cursor.carry.size / sample_rate - chunk_media) > 0.001:
492 # the leftover belongs to audio we are no longer continuing from
493 cursor.carry = np.zeros(0, dtype=np.float32)
494 cursor.carry_media = chunk_media
495 mono = np.concatenate([cursor.carry, mono])
496 offset = 0
497 while mono.size - offset >= WAVE_SAMPLES:
498 window = mono[offset : offset + WAVE_SAMPLES]
499 offset += WAVE_SAMPLES
500 quantized = np.rint(np.clip(window, -1.0, 1.0) * 127.0 + 128.0).astype(np.uint8)
501 # stamped at the end of the window, the instant it finishes sounding
502 end_media = cursor.carry_media + offset / sample_rate
503 frame = pack_wave_frame(cursor.media_to_clock_us(end_media), quantized.tobytes())
504 tap.ring.append(frame)
505 tap.fan_out(frame)
506 cursor.carry = mono[offset:].copy()
507 cursor.carry_media += offset / sample_rate
508 cursor.next_chunk += 1
509
510 def _sync_color(self, tap: Tap) -> None:
511 """Fan out the track palette whenever it changes (once per track, in practice)."""
512 if not self.provider.config.get_value(CONF_COLOR_TINT):
513 return
514 player = self.mass.players.get_player(tap.player_id)
515 media = player.state.current_media if player is not None else None
516 payload = palette_payload(media.palette if media is not None else None)
517 if payload != tap.last_color:
518 tap.apply_color(payload)
519
520 def _schedule_beats(
521 self, tap: Tap, item: QueueItem, anchor_us: int, speed: float = 1.0
522 ) -> None:
523 """
524 (Re)build a tap's beat schedule for a track.
525
526 A re-anchor of an item whose analysis is already cached (a seek)
527 rebuilds the schedule in place, without a task or a new query.
528
529 :param tap: The tap to fan the beats out to.
530 :param item: The queue item now playing.
531 :param anchor_us: Clock time of that item's media time zero.
532 :param speed: Playback speed the queue plays the item at.
533 """
534 if tap.beats_analysis is not None and tap.beats_analysis[0] == item.queue_item_id:
535 self._fan_out_beats(tap, tap.beats_analysis[1], anchor_us, speed)
536 return
537 tap.beats_task = self.mass.create_task(
538 self._hydrate_beats(tap, item, anchor_us, speed),
539 task_id=f"milkdrop_beats_{tap.player_id}",
540 abort_existing=True,
541 )
542
543 async def _hydrate_beats(
544 self, tap: Tap, item: QueueItem, anchor_us: int, speed: float = 1.0
545 ) -> None:
546 """Wait for a track's beat analysis and schedule the beats that are still ahead."""
547 streamdetails = item.streamdetails
548 if streamdetails is None:
549 return
550 # smart_fades is the only AA provider that produces beats. Without it,
551 # no amount of waiting will turn any up.
552 if not self.mass.streams.audio_analysis.smart_fades_provider_available:
553 return
554 for _ in range(BEAT_RETRY_ATTEMPTS):
555 analysis = await self.mass.streams.audio_analysis.get_audio_analysis(
556 streamdetails.item_id,
557 streamdetails.provider,
558 media_type=streamdetails.media_type,
559 priority=(SMART_FADES_ANALYSIS_DOMAIN,),
560 )
561 if analysis is not None and analysis.beats:
562 break
563 await asyncio.sleep(BEAT_RETRY_SECONDS)
564 else:
565 self.logger.debug("No beat analysis for %s", streamdetails.uri)
566 return
567 tap.beats_analysis = (item.queue_item_id, analysis)
568 self._fan_out_beats(tap, analysis, anchor_us, speed)
569
570 def _fan_out_beats(
571 self, tap: Tap, analysis: AudioAnalysisData, anchor_us: int, speed: float = 1.0
572 ) -> None:
573 """
574 Schedule the analysis beats that are still ahead and fan them out.
575
576 :param tap: The tap to fan the beats out to.
577 :param analysis: The beat analysis of the item now playing.
578 :param anchor_us: Clock time of that item's media time zero.
579 :param speed: Playback speed the queue plays the item at.
580 """
581 beats = analysis.beats or ()
582 downbeats = {float(value) for value in analysis.downbeats or ()}
583 now_us = server_now_us()
584 scheduled = 0
585 for beat in beats:
586 timestamp_us = anchor_us + int(float(beat) / speed * 1_000_000)
587 if timestamp_us <= now_us:
588 continue
589 frame = pack_beat_frame(timestamp_us, float(beat) in downbeats)
590 tap.beats.append((timestamp_us, frame))
591 tap.fan_out(frame)
592 scheduled += 1
593 self.logger.debug("Scheduled %s of %s beat(s)", scheduled, len(beats))
594