1144 lines
41 KiB
Rust
1144 lines
41 KiB
Rust
//! Local, offline-first music embeddings and exact cosine search.
|
||
//!
|
||
//! SQLite owns the durable vectors. The in-memory index is deliberately
|
||
//! replaceable: it is rebuilt for the active profile and never becomes a
|
||
//! second source of truth.
|
||
|
||
use std::collections::{HashMap, VecDeque};
|
||
use std::fs::File;
|
||
use std::path::{Path, PathBuf};
|
||
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
|
||
use std::sync::{Arc, Mutex, RwLock};
|
||
use std::time::Instant;
|
||
|
||
use anyhow::{Context as _, Result};
|
||
use futures_util::StreamExt as _;
|
||
use rodio::{Decoder, Source as _};
|
||
use rustfft::FftPlanner;
|
||
use rustfft::num_complex::Complex;
|
||
use sha2::{Digest as _, Sha256};
|
||
use tokio::io::AsyncWriteExt as _;
|
||
use tract_onnx::prelude::*;
|
||
use tract_onnx::tract_core::dims;
|
||
|
||
use crate::app::event::AppEvent;
|
||
use crate::config::settings::SimilaritySettings;
|
||
use crate::library::models::TrackItem;
|
||
use crate::library::{Library, StoredEmbedding};
|
||
|
||
const SAMPLE_RATE: usize = 16_000;
|
||
const FRAME_SIZE: usize = 512;
|
||
const HOP_SIZE: usize = 256;
|
||
const MEL_BANDS: usize = 96;
|
||
const PATCH_FRAMES: usize = 128;
|
||
const PATCH_HOP: usize = 62;
|
||
const EMBEDDING_DIMENSIONS: usize = 1280;
|
||
const MODEL_BATCH: usize = 8;
|
||
const MAX_MODEL_BYTES: usize = 64 * 1024 * 1024;
|
||
const RESULT_LIMIT: usize = 50;
|
||
const MAX_PER_ARTIST: usize = 3;
|
||
const NEAR_DUPLICATE_COSINE: f32 = 0.995;
|
||
const FULL_TRACK_MAX_SECONDS: u32 = 5 * 60;
|
||
const LONG_TRACK_WINDOW_SECONDS: u32 = 60;
|
||
const LONG_TRACK_WINDOWS: usize = 3;
|
||
|
||
pub const DEFAULT_MODEL_ID: &str = "discogs-effnet-bsdynamic-1";
|
||
pub const DEFAULT_PROFILE_ID: &str = "furumi-full-track-v1";
|
||
|
||
#[derive(Debug, Clone, Copy)]
|
||
pub struct ProfileSpec {
|
||
pub id: &'static str,
|
||
pub title: &'static str,
|
||
}
|
||
|
||
pub const PROFILES: &[ProfileSpec] = &[ProfileSpec {
|
||
id: DEFAULT_PROFILE_ID,
|
||
title: "Full track / balanced long track",
|
||
}];
|
||
|
||
#[derive(Debug, Clone, Copy)]
|
||
pub struct ModelSpec {
|
||
pub id: &'static str,
|
||
pub version: &'static str,
|
||
pub filename: &'static str,
|
||
pub url: &'static str,
|
||
pub sha256: &'static str,
|
||
pub dimensions: usize,
|
||
pub license: &'static str,
|
||
}
|
||
|
||
pub const MODELS: &[ModelSpec] = &[ModelSpec {
|
||
id: DEFAULT_MODEL_ID,
|
||
version: "1",
|
||
filename: "discogs-effnet-bsdynamic-1.onnx",
|
||
url: "https://essentia.upf.edu/models/feature-extractors/discogs-effnet/discogs-effnet-bsdynamic-1.onnx",
|
||
sha256: "a280825b334797cf677939db8cd5762c0392aedd0ca6415dbc1cd083f045e43c",
|
||
dimensions: EMBEDDING_DIMENSIONS,
|
||
license: "CC BY-NC-SA 4.0 (or proprietary from MTG)",
|
||
}];
|
||
|
||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||
pub enum Phase {
|
||
#[default]
|
||
Disabled,
|
||
Downloading,
|
||
Loading,
|
||
Processing,
|
||
Ready,
|
||
Error,
|
||
}
|
||
|
||
impl Phase {
|
||
pub fn label(self) -> &'static str {
|
||
match self {
|
||
Self::Disabled => "disabled",
|
||
Self::Downloading => "downloading model",
|
||
Self::Loading => "loading model",
|
||
Self::Processing => "processing",
|
||
Self::Ready => "ready",
|
||
Self::Error => "error",
|
||
}
|
||
}
|
||
}
|
||
|
||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||
pub struct SimilarityStatus {
|
||
pub phase: Phase,
|
||
pub active_profile: Option<String>,
|
||
pub target_profile: Option<String>,
|
||
pub model: String,
|
||
pub total_tracks: usize,
|
||
pub completed_tracks: usize,
|
||
pub failed_tracks: usize,
|
||
pub stored_vectors: usize,
|
||
pub stored_bytes: u64,
|
||
pub current_track: Option<String>,
|
||
pub last_error: Option<String>,
|
||
}
|
||
|
||
#[derive(Debug, Clone)]
|
||
pub struct QueryVector {
|
||
pub profile_id: String,
|
||
pub vector: Vec<f32>,
|
||
pub source_content_id: Option<String>,
|
||
}
|
||
|
||
#[derive(Debug, Clone)]
|
||
pub struct SimilarTrack {
|
||
pub track: TrackItem,
|
||
pub score: f32,
|
||
pub embedding_signature: [u8; music_dht::similarity::SIMILARITY_SIGNATURE_BYTES],
|
||
}
|
||
|
||
#[derive(Default)]
|
||
struct Index {
|
||
profile_id: Option<String>,
|
||
entries: Vec<StoredEmbedding>,
|
||
}
|
||
|
||
type RunnableModel = Arc<TypedRunnableModel>;
|
||
|
||
pub struct Manager {
|
||
library: Arc<Library>,
|
||
event_tx: tokio::sync::mpsc::UnboundedSender<AppEvent>,
|
||
settings: Mutex<SimilaritySettings>,
|
||
workers: AtomicUsize,
|
||
generation: AtomicU64,
|
||
pipeline_running: AtomicBool,
|
||
rescan_requested: AtomicBool,
|
||
status: Mutex<SimilarityStatus>,
|
||
index: RwLock<Index>,
|
||
model: Mutex<Option<(String, RunnableModel)>>,
|
||
model_dir: PathBuf,
|
||
}
|
||
|
||
impl Manager {
|
||
pub fn new(
|
||
library: Arc<Library>,
|
||
event_tx: tokio::sync::mpsc::UnboundedSender<AppEvent>,
|
||
settings: SimilaritySettings,
|
||
) -> Arc<Self> {
|
||
let model_dir = crate::config::project_dirs()
|
||
.map(|dirs| dirs.cache_dir().join("similarity-models"))
|
||
.unwrap_or_else(|| PathBuf::from("similarity-models"));
|
||
let mut index = Index::default();
|
||
if let Some(profile_id) = settings.active_profile.as_deref() {
|
||
match library.load_similarity_index(profile_id) {
|
||
Ok(entries) => {
|
||
index.profile_id = Some(profile_id.to_string());
|
||
index.entries = entries;
|
||
}
|
||
Err(err) => tracing::warn!(%err, "similarity index restore failed"),
|
||
}
|
||
}
|
||
let target_profile = model_by_id(&settings.model)
|
||
.filter(|_| profile_by_id(&settings.profile).is_some())
|
||
.map(|model| profile_fingerprint(model, &settings.profile));
|
||
let restored_profile_is_current = index.profile_id == target_profile;
|
||
let status = SimilarityStatus {
|
||
phase: if !settings.enabled {
|
||
Phase::Disabled
|
||
} else if restored_profile_is_current {
|
||
Phase::Ready
|
||
} else {
|
||
Phase::Loading
|
||
},
|
||
active_profile: index.profile_id.clone(),
|
||
target_profile,
|
||
model: settings.model.clone(),
|
||
..SimilarityStatus::default()
|
||
};
|
||
Arc::new(Self {
|
||
library,
|
||
event_tx,
|
||
workers: AtomicUsize::new(settings.workers.clamp(1, 16)),
|
||
generation: AtomicU64::new(0),
|
||
pipeline_running: AtomicBool::new(false),
|
||
rescan_requested: AtomicBool::new(false),
|
||
settings: Mutex::new(settings),
|
||
status: Mutex::new(status),
|
||
index: RwLock::new(index),
|
||
model: Mutex::new(None),
|
||
model_dir,
|
||
})
|
||
}
|
||
|
||
pub fn settings(&self) -> SimilaritySettings {
|
||
lock(&self.settings).clone()
|
||
}
|
||
|
||
pub fn status(&self) -> SimilarityStatus {
|
||
lock(&self.status).clone()
|
||
}
|
||
|
||
pub fn network_allowed(&self) -> bool {
|
||
let settings = lock(&self.settings);
|
||
settings.enabled && settings.federation_consent
|
||
}
|
||
|
||
pub fn apply(self: &Arc<Self>, settings: SimilaritySettings) {
|
||
self.workers
|
||
.store(settings.workers.clamp(1, 16), Ordering::Release);
|
||
let previous = std::mem::replace(&mut *lock(&self.settings), settings.clone());
|
||
if !settings.enabled {
|
||
self.generation.fetch_add(1, Ordering::AcqRel);
|
||
self.update_status(|status| {
|
||
status.phase = Phase::Disabled;
|
||
status.current_track = None;
|
||
status.target_profile = None;
|
||
status.last_error = None;
|
||
});
|
||
return;
|
||
}
|
||
if !previous.enabled
|
||
|| previous.model != settings.model
|
||
|| previous.profile != settings.profile
|
||
{
|
||
self.generation.fetch_add(1, Ordering::AcqRel);
|
||
self.start();
|
||
}
|
||
}
|
||
|
||
/// Requests a scan without cancelling useful work already in progress.
|
||
/// Bursts of library-change notifications collapse into one follow-up
|
||
/// pass, so metadata refreshes cannot repeatedly restart the model.
|
||
pub fn start(self: &Arc<Self>) {
|
||
self.rescan_requested.store(true, Ordering::Release);
|
||
if self.pipeline_running.swap(true, Ordering::AcqRel) {
|
||
return;
|
||
}
|
||
let this = Arc::clone(self);
|
||
tokio::spawn(async move {
|
||
loop {
|
||
// This pass covers every notification received before it
|
||
// starts. A notification during the pass requests one more.
|
||
this.rescan_requested.store(false, Ordering::Release);
|
||
let generation = this.generation.load(Ordering::Acquire);
|
||
if let Err(err) = this.run_pipeline(generation).await
|
||
&& this.generation.load(Ordering::Acquire) == generation
|
||
{
|
||
tracing::error!(%err, "similarity pipeline failed");
|
||
this.update_status(|status| {
|
||
status.phase = Phase::Error;
|
||
status.current_track = None;
|
||
status.last_error = Some(format!("{err:#}"));
|
||
});
|
||
}
|
||
|
||
if this.rescan_requested.load(Ordering::Acquire) {
|
||
continue;
|
||
}
|
||
|
||
this.pipeline_running.store(false, Ordering::Release);
|
||
// Close the small race between checking the request flag and
|
||
// releasing ownership of the worker. If another worker has
|
||
// already claimed it, that worker owns the pending pass.
|
||
if this.rescan_requested.swap(false, Ordering::AcqRel)
|
||
&& !this.pipeline_running.swap(true, Ordering::AcqRel)
|
||
{
|
||
continue;
|
||
}
|
||
break;
|
||
}
|
||
});
|
||
}
|
||
|
||
pub fn clear(self: &Arc<Self>) {
|
||
self.generation.fetch_add(1, Ordering::AcqRel);
|
||
let this = Arc::clone(self);
|
||
tokio::spawn(async move {
|
||
let library = Arc::clone(&this.library);
|
||
let result = tokio::task::spawn_blocking(move || library.clear_similarity_embeddings())
|
||
.await
|
||
.context("embedding clear task failed")
|
||
.and_then(|result| result);
|
||
match result {
|
||
Ok(()) => {
|
||
*write(&this.index) = Index::default();
|
||
lock(&this.settings).active_profile = None;
|
||
this.update_status(|status| {
|
||
*status = SimilarityStatus {
|
||
phase: if this.settings().enabled {
|
||
Phase::Loading
|
||
} else {
|
||
Phase::Disabled
|
||
},
|
||
model: this.settings().model,
|
||
..SimilarityStatus::default()
|
||
};
|
||
});
|
||
let _ = this
|
||
.event_tx
|
||
.send(AppEvent::SimilarityProfileActivated(None));
|
||
if this.settings().enabled {
|
||
this.start();
|
||
}
|
||
}
|
||
Err(err) => this.update_status(|status| {
|
||
status.phase = Phase::Error;
|
||
status.last_error = Some(format!("clear failed: {err:#}"));
|
||
}),
|
||
}
|
||
});
|
||
}
|
||
|
||
pub fn query_for_track(&self, track_id: i64) -> Result<QueryVector> {
|
||
let profile_id = read(&self.index)
|
||
.profile_id
|
||
.clone()
|
||
.context("no similarity profile is ready yet")?;
|
||
let vector = self
|
||
.library
|
||
.similarity_embedding(track_id, &profile_id)?
|
||
.context("this track has not been processed yet")?;
|
||
let source_track = self.library.tracks_by_ids(&[track_id])?.into_iter().next();
|
||
let source_content_id = source_track.as_ref().and_then(|track| {
|
||
track
|
||
.content_id
|
||
.clone()
|
||
.or_else(|| crate::library::audio_content_id(&track.file_path))
|
||
});
|
||
Ok(QueryVector {
|
||
profile_id,
|
||
vector,
|
||
source_content_id,
|
||
})
|
||
}
|
||
|
||
pub fn search_track(
|
||
&self,
|
||
track_id: i64,
|
||
limit: usize,
|
||
) -> Result<(Vec<SimilarTrack>, QueryVector)> {
|
||
let query = self.query_for_track(track_id)?;
|
||
let matches = self.search_vector(
|
||
&query.profile_id,
|
||
&query.vector,
|
||
Some(track_id),
|
||
query.source_content_id.as_deref(),
|
||
limit,
|
||
)?;
|
||
Ok((matches, query))
|
||
}
|
||
|
||
pub fn search_vector(
|
||
&self,
|
||
profile_id: &str,
|
||
vector: &[f32],
|
||
exclude_track_id: Option<i64>,
|
||
exclude_content_id: Option<&str>,
|
||
limit: usize,
|
||
) -> Result<Vec<SimilarTrack>> {
|
||
anyhow::ensure!(
|
||
!vector.is_empty() && vector.len() <= 4096,
|
||
"wrong embedding dimensions"
|
||
);
|
||
anyhow::ensure!(
|
||
vector.iter().all(|value| value.is_finite()),
|
||
"invalid embedding"
|
||
);
|
||
let index = read(&self.index);
|
||
anyhow::ensure!(
|
||
index.profile_id.as_deref() == Some(profile_id),
|
||
"the requested similarity profile is not active"
|
||
);
|
||
let mut scores: Vec<(i64, f32, &str, &[f32])> = index
|
||
.entries
|
||
.iter()
|
||
.filter(|entry| {
|
||
Some(entry.track_id) != exclude_track_id && entry.vector.len() == vector.len()
|
||
})
|
||
.map(|entry| {
|
||
(
|
||
entry.track_id,
|
||
dot(vector, &entry.vector),
|
||
entry.artist_key.as_str(),
|
||
entry.vector.as_slice(),
|
||
)
|
||
})
|
||
.filter(|(_, score, _, _)| score.is_finite())
|
||
.collect();
|
||
scores.sort_by(|left, right| right.1.total_cmp(&left.1));
|
||
|
||
// Pull a wider candidate set, then cap each primary artist so a large
|
||
// discography cannot fill the whole result page.
|
||
let mut artist_counts: HashMap<String, usize> = HashMap::new();
|
||
let mut kept_vectors = vec![vector];
|
||
let mut selected = Vec::new();
|
||
for (track_id, score, artist, candidate_vector) in scores {
|
||
if is_near_duplicate(candidate_vector, &kept_vectors) {
|
||
continue;
|
||
}
|
||
let count = artist_counts.entry(artist.to_string()).or_default();
|
||
if !artist.is_empty() && *count >= MAX_PER_ARTIST {
|
||
continue;
|
||
}
|
||
*count += 1;
|
||
let embedding_signature = music_dht::similarity::embedding_signature(candidate_vector)?;
|
||
kept_vectors.push(candidate_vector);
|
||
selected.push((track_id, score, embedding_signature));
|
||
if selected.len() >= limit.min(RESULT_LIMIT) {
|
||
break;
|
||
}
|
||
}
|
||
drop(index);
|
||
|
||
let ids: Vec<i64> = selected.iter().map(|(id, _, _)| *id).collect();
|
||
let tracks = self.library.tracks_by_ids(&ids)?;
|
||
let by_id: HashMap<i64, TrackItem> =
|
||
tracks.into_iter().map(|track| (track.id, track)).collect();
|
||
Ok(selected
|
||
.into_iter()
|
||
.filter_map(|(id, score, signature)| {
|
||
by_id
|
||
.get(&id)
|
||
.cloned()
|
||
.map(|track| (track, score, signature))
|
||
})
|
||
.filter(|(track, _, _)| {
|
||
!exclude_content_id
|
||
.is_some_and(|source| track.content_id.as_deref() == Some(source))
|
||
})
|
||
.map(|(track, score, embedding_signature)| SimilarTrack {
|
||
track,
|
||
score,
|
||
embedding_signature,
|
||
})
|
||
.collect())
|
||
}
|
||
|
||
async fn run_pipeline(self: &Arc<Self>, generation: u64) -> Result<()> {
|
||
let settings = self.settings();
|
||
if !settings.enabled {
|
||
return Ok(());
|
||
}
|
||
let spec = model_by_id(&settings.model)
|
||
.with_context(|| format!("unknown similarity model '{}'", settings.model))?;
|
||
anyhow::ensure!(
|
||
profile_by_id(&settings.profile).is_some(),
|
||
"unknown preprocessing profile '{}'",
|
||
settings.profile
|
||
);
|
||
let profile_id = profile_fingerprint(spec, &settings.profile);
|
||
self.library.ensure_similarity_profile(
|
||
&profile_id,
|
||
spec.id,
|
||
spec.version,
|
||
spec.sha256,
|
||
&settings.profile,
|
||
spec.dimensions,
|
||
)?;
|
||
let stats = self.library.similarity_storage_stats(&profile_id)?;
|
||
self.update_status(|status| {
|
||
status.target_profile = Some(profile_id.clone());
|
||
status.model = spec.id.to_string();
|
||
status.total_tracks = stats.total_tracks;
|
||
status.completed_tracks = stats.embedded_tracks;
|
||
status.stored_vectors = stats.stored_vectors;
|
||
status.stored_bytes = stats.stored_bytes;
|
||
status.failed_tracks = 0;
|
||
status.current_track = None;
|
||
status.last_error = None;
|
||
});
|
||
|
||
let mut pending: VecDeque<_> = self.library.pending_similarity_tracks(&profile_id)?.into();
|
||
if pending.is_empty() {
|
||
self.ensure_generation(generation)?;
|
||
return self.activate_profile(profile_id);
|
||
}
|
||
|
||
let model_path = self.ensure_model(spec, generation).await?;
|
||
self.ensure_generation(generation)?;
|
||
let model = self.load_model(&profile_id, &model_path).await?;
|
||
self.ensure_generation(generation)?;
|
||
self.update_status(|status| status.phase = Phase::Processing);
|
||
|
||
let mut jobs = tokio::task::JoinSet::new();
|
||
while !pending.is_empty() || !jobs.is_empty() {
|
||
self.ensure_generation(generation)?;
|
||
let workers = self.workers.load(Ordering::Acquire).clamp(1, 16);
|
||
while jobs.len() < workers {
|
||
let Some(track) = pending.pop_front() else {
|
||
break;
|
||
};
|
||
self.update_status(|status| status.current_track = Some(track.title.clone()));
|
||
let library = Arc::clone(&self.library);
|
||
let model = Arc::clone(&model);
|
||
let profile_id = profile_id.clone();
|
||
jobs.spawn_blocking(move || {
|
||
let started = Instant::now();
|
||
let result =
|
||
embed_track(&model, Path::new(&track.file_path), track.duration_seconds)
|
||
.and_then(|vector| {
|
||
library.store_similarity_embedding(&track, &profile_id, &vector)
|
||
});
|
||
(track, result, started.elapsed())
|
||
});
|
||
}
|
||
let Some(result) = jobs.join_next().await else {
|
||
continue;
|
||
};
|
||
let (track, result, elapsed) = result.context("embedding worker panicked")?;
|
||
self.ensure_generation(generation)?;
|
||
match result {
|
||
Ok(()) => {
|
||
tracing::info!(
|
||
track_id = track.id,
|
||
title = %track.title,
|
||
elapsed_ms = elapsed.as_millis(),
|
||
profile = %profile_id,
|
||
"track embedding calculated"
|
||
);
|
||
self.update_status(|status| status.completed_tracks += 1);
|
||
}
|
||
Err(err) => {
|
||
tracing::warn!(track_id = track.id, title = %track.title, %err, "track embedding failed");
|
||
self.update_status(|status| {
|
||
status.failed_tracks += 1;
|
||
status.last_error = Some(format!("{}: {err:#}", track.title));
|
||
});
|
||
}
|
||
}
|
||
}
|
||
self.ensure_generation(generation)?;
|
||
self.activate_profile(profile_id)
|
||
}
|
||
|
||
fn activate_profile(&self, profile_id: String) -> Result<()> {
|
||
let entries = self.library.load_similarity_index(&profile_id)?;
|
||
let total_tracks = self
|
||
.library
|
||
.similarity_storage_stats(&profile_id)?
|
||
.total_tracks;
|
||
anyhow::ensure!(
|
||
total_tracks == 0 || !entries.is_empty(),
|
||
"no tracks could be processed"
|
||
);
|
||
*write(&self.index) = Index {
|
||
profile_id: Some(profile_id.clone()),
|
||
entries,
|
||
};
|
||
let profile_changed = {
|
||
let mut settings = lock(&self.settings);
|
||
let changed = settings.active_profile.as_deref() != Some(&profile_id);
|
||
settings.active_profile = Some(profile_id.clone());
|
||
changed
|
||
};
|
||
let stats = self.library.similarity_storage_stats(&profile_id)?;
|
||
self.update_status(|status| {
|
||
status.phase = Phase::Ready;
|
||
status.active_profile = Some(profile_id.clone());
|
||
status.target_profile = Some(profile_id.clone());
|
||
status.total_tracks = stats.total_tracks;
|
||
status.completed_tracks = stats.embedded_tracks;
|
||
status.stored_vectors = stats.stored_vectors;
|
||
status.stored_bytes = stats.stored_bytes;
|
||
status.current_track = None;
|
||
});
|
||
if profile_changed {
|
||
let _ = self
|
||
.event_tx
|
||
.send(AppEvent::SimilarityProfileActivated(Some(profile_id)));
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
fn ensure_generation(&self, generation: u64) -> Result<()> {
|
||
anyhow::ensure!(
|
||
self.generation.load(Ordering::Acquire) == generation,
|
||
"similarity processing superseded by newer settings"
|
||
);
|
||
Ok(())
|
||
}
|
||
|
||
async fn ensure_model(&self, spec: &ModelSpec, generation: u64) -> Result<PathBuf> {
|
||
tokio::fs::create_dir_all(&self.model_dir).await?;
|
||
let path = self.model_dir.join(spec.filename);
|
||
if path.exists() {
|
||
let verify_path = path.clone();
|
||
let expected = spec.sha256.to_string();
|
||
let valid = tokio::task::spawn_blocking(move || sha256_file(&verify_path))
|
||
.await
|
||
.context("model hash task failed")??
|
||
== expected;
|
||
if valid {
|
||
return Ok(path);
|
||
}
|
||
tokio::fs::remove_file(&path).await?;
|
||
}
|
||
|
||
self.update_status(|status| status.phase = Phase::Downloading);
|
||
let response = reqwest::get(spec.url).await?.error_for_status()?;
|
||
let tmp = path.with_extension(format!("part-{}-{generation}", std::process::id()));
|
||
let mut file = tokio::fs::File::create(&tmp).await?;
|
||
let mut hasher = Sha256::new();
|
||
let mut received = 0usize;
|
||
let mut stream = response.bytes_stream();
|
||
while let Some(chunk) = stream.next().await {
|
||
let chunk = chunk?;
|
||
received = received.saturating_add(chunk.len());
|
||
anyhow::ensure!(
|
||
received <= MAX_MODEL_BYTES,
|
||
"model download exceeds size limit"
|
||
);
|
||
hasher.update(&chunk);
|
||
file.write_all(&chunk).await?;
|
||
}
|
||
file.flush().await?;
|
||
drop(file);
|
||
let actual = format!("{:x}", hasher.finalize());
|
||
if actual != spec.sha256 {
|
||
let _ = tokio::fs::remove_file(&tmp).await;
|
||
anyhow::bail!("downloaded model hash mismatch");
|
||
}
|
||
if let Err(err) = tokio::fs::rename(&tmp, &path).await {
|
||
// A superseding pipeline may have installed the same verified
|
||
// artifact first. This is expected on platforms where rename
|
||
// does not replace an existing destination.
|
||
if path.exists() {
|
||
let _ = tokio::fs::remove_file(&tmp).await;
|
||
} else {
|
||
return Err(err.into());
|
||
}
|
||
}
|
||
Ok(path)
|
||
}
|
||
|
||
async fn load_model(&self, profile_id: &str, path: &Path) -> Result<RunnableModel> {
|
||
if let Some((cached_profile, model)) = lock(&self.model).as_ref()
|
||
&& cached_profile == profile_id
|
||
{
|
||
return Ok(Arc::clone(model));
|
||
}
|
||
self.update_status(|status| status.phase = Phase::Loading);
|
||
let path = path.to_path_buf();
|
||
let model = tokio::task::spawn_blocking(move || load_onnx(&path))
|
||
.await
|
||
.context("model loading task failed")??;
|
||
*lock(&self.model) = Some((profile_id.to_string(), Arc::clone(&model)));
|
||
Ok(model)
|
||
}
|
||
|
||
fn update_status(&self, update: impl FnOnce(&mut SimilarityStatus)) {
|
||
let snapshot = {
|
||
let mut status = lock(&self.status);
|
||
let previous = status.clone();
|
||
update(&mut status);
|
||
(*status != previous).then(|| status.clone())
|
||
};
|
||
if let Some(snapshot) = snapshot {
|
||
let _ = self.event_tx.send(AppEvent::SimilarityStatus(snapshot));
|
||
}
|
||
}
|
||
}
|
||
|
||
pub fn model_by_id(id: &str) -> Option<&'static ModelSpec> {
|
||
MODELS.iter().find(|model| model.id == id)
|
||
}
|
||
|
||
pub fn profile_by_id(id: &str) -> Option<&'static ProfileSpec> {
|
||
PROFILES.iter().find(|profile| profile.id == id)
|
||
}
|
||
|
||
pub fn profile_details(profile_id: &str, model_id: &str) -> Option<String> {
|
||
let profile = profile_by_id(profile_id)?;
|
||
let dimensions = model_by_id(model_id)
|
||
.map(|model| model.dimensions.to_string())
|
||
.unwrap_or_else(|| "model-defined".to_string());
|
||
Some(format!(
|
||
"{}\n\nTrack selection:\n• Up to {} seconds: entire track.\n• Longer: {} × {}-second windows (start, middle, end).\n\nAudio: mono, {} Hz; 16-tap windowed-sinc resampling.\nSpectrogram: Hann window; FFT {}, hop {} samples (16 ms).\nMel: {} Slaney bands, 0–8 kHz, unit-triangle normalization.\nCompression: log10(1 + 10000 × energy).\nPatches: {} frames (~2.05 s), hop {} frames (~0.99 s).\nAggregation: mean of patch embeddings, then L2 normalization.\nOutput dimensions for selected model: {}.\n\nCompatibility includes the exact model version and SHA-256; peers compare only matching profiles.\n\nEnter / Esc: close",
|
||
profile.title,
|
||
FULL_TRACK_MAX_SECONDS,
|
||
LONG_TRACK_WINDOWS,
|
||
LONG_TRACK_WINDOW_SECONDS,
|
||
SAMPLE_RATE,
|
||
FRAME_SIZE,
|
||
HOP_SIZE,
|
||
MEL_BANDS,
|
||
PATCH_FRAMES,
|
||
PATCH_HOP,
|
||
dimensions,
|
||
))
|
||
}
|
||
|
||
pub fn profile_fingerprint(model: &ModelSpec, profile: &str) -> String {
|
||
let contract = format!(
|
||
"furumi-similarity-v1\nmodel={}\nversion={}\nsha256={}\nprofile={}\ninput={}-mono-windowed-sinc16\nselection=full-to-{}s-else-first-middle-last-{}s\nframe={}\nhop={}\nmel=slaney-{}-unit-tri\npatch={}\npatch-hop={}\naggregate=mean-l2\ndimensions={}",
|
||
model.id,
|
||
model.version,
|
||
model.sha256,
|
||
profile,
|
||
SAMPLE_RATE,
|
||
FULL_TRACK_MAX_SECONDS,
|
||
LONG_TRACK_WINDOW_SECONDS,
|
||
FRAME_SIZE,
|
||
HOP_SIZE,
|
||
MEL_BANDS,
|
||
PATCH_FRAMES,
|
||
PATCH_HOP,
|
||
model.dimensions
|
||
);
|
||
format!("sim1:{}", blake3::hash(contract.as_bytes()).to_hex())
|
||
}
|
||
|
||
fn load_onnx(path: &Path) -> Result<RunnableModel> {
|
||
let model = tract_onnx::onnx().model_for_path(path)?;
|
||
let batch = model.sym("batch_size");
|
||
let model = model
|
||
.with_input_fact(0, f32::fact(dims!(batch, PATCH_FRAMES, MEL_BANDS)).into())?
|
||
.into_optimized()?
|
||
.into_runnable()?;
|
||
Ok(model)
|
||
}
|
||
|
||
fn embed_track(model: &RunnableModel, path: &Path, duration_seconds: f64) -> Result<Vec<f32>> {
|
||
let signal = decode_mono_16k(path, duration_seconds)?;
|
||
let mel = mel_spectrogram(&signal)?;
|
||
anyhow::ensure!(
|
||
mel.len() >= PATCH_FRAMES,
|
||
"track is too short for the model"
|
||
);
|
||
let starts: Vec<usize> = (0..=mel.len() - PATCH_FRAMES).step_by(PATCH_HOP).collect();
|
||
let mut sum = vec![0.0f32; EMBEDDING_DIMENSIONS];
|
||
let mut count = 0usize;
|
||
for batch in starts.chunks(MODEL_BATCH) {
|
||
let mut input = vec![0.0f32; MODEL_BATCH * PATCH_FRAMES * MEL_BANDS];
|
||
for (batch_index, &start) in batch.iter().enumerate() {
|
||
let offset = batch_index * PATCH_FRAMES * MEL_BANDS;
|
||
for frame in 0..PATCH_FRAMES {
|
||
let dst = offset + frame * MEL_BANDS;
|
||
input[dst..dst + MEL_BANDS].copy_from_slice(&mel[start + frame]);
|
||
}
|
||
}
|
||
let tensor = Tensor::from_shape(&[MODEL_BATCH, PATCH_FRAMES, MEL_BANDS], &input)?;
|
||
let outputs = model.run(tvec!(tensor.into_tvalue()))?;
|
||
let embedding = outputs
|
||
.iter()
|
||
.find(|output| output.len() == MODEL_BATCH * EMBEDDING_DIMENSIONS)
|
||
.context("model did not return its 1280-dimensional embedding output")?
|
||
.to_plain_array_view::<f32>()?;
|
||
let values = embedding
|
||
.as_slice()
|
||
.context("model embedding output is not contiguous")?;
|
||
for batch_index in 0..batch.len() {
|
||
let row = &values
|
||
[batch_index * EMBEDDING_DIMENSIONS..(batch_index + 1) * EMBEDDING_DIMENSIONS];
|
||
for (total, value) in sum.iter_mut().zip(row) {
|
||
*total += *value;
|
||
}
|
||
count += 1;
|
||
}
|
||
}
|
||
anyhow::ensure!(count > 0, "model produced no patches");
|
||
for value in &mut sum {
|
||
*value /= count as f32;
|
||
}
|
||
normalize(&mut sum)?;
|
||
Ok(sum)
|
||
}
|
||
|
||
fn decode_mono_16k(path: &Path, duration_seconds: f64) -> Result<Vec<f32>> {
|
||
if duration_seconds.is_finite() && duration_seconds > f64::from(FULL_TRACK_MAX_SECONDS) {
|
||
let window = f64::from(LONG_TRACK_WINDOW_SECONDS);
|
||
let starts = [
|
||
0.0,
|
||
(duration_seconds / 2.0 - window / 2.0).max(0.0),
|
||
(duration_seconds - window).max(0.0),
|
||
];
|
||
let mut selected = Vec::new();
|
||
for start in starts {
|
||
selected.extend(decode_mono_window(path, start, Some(window))?);
|
||
}
|
||
anyhow::ensure!(!selected.is_empty(), "decoded track is empty");
|
||
return Ok(selected);
|
||
}
|
||
decode_mono_window(path, 0.0, None)
|
||
}
|
||
|
||
fn decode_mono_window(
|
||
path: &Path,
|
||
start_seconds: f64,
|
||
length_seconds: Option<f64>,
|
||
) -> Result<Vec<f32>> {
|
||
let file = File::open(path).with_context(|| format!("opening {}", path.display()))?;
|
||
let mut decoder =
|
||
Decoder::try_from(file).with_context(|| format!("decoding {}", path.display()))?;
|
||
let channels = decoder.channels().get() as usize;
|
||
let source_rate = decoder.sample_rate().get() as usize;
|
||
if start_seconds > 0.0 {
|
||
decoder
|
||
.try_seek(std::time::Duration::from_secs_f64(start_seconds))
|
||
.with_context(|| format!("seeking {}", path.display()))?;
|
||
}
|
||
let max_samples =
|
||
length_seconds.map(|seconds| (seconds * source_rate as f64).ceil() as usize * channels);
|
||
let mut mono = Vec::new();
|
||
let mut channel_sum = 0.0f32;
|
||
let mut channel_index = 0usize;
|
||
for (sample_index, sample) in decoder.enumerate() {
|
||
if max_samples.is_some_and(|limit| sample_index >= limit) {
|
||
break;
|
||
}
|
||
channel_sum += sample;
|
||
channel_index += 1;
|
||
if channel_index == channels {
|
||
mono.push(channel_sum / channels as f32);
|
||
channel_sum = 0.0;
|
||
channel_index = 0;
|
||
}
|
||
}
|
||
anyhow::ensure!(!mono.is_empty(), "decoded track is empty");
|
||
if source_rate == SAMPLE_RATE {
|
||
return Ok(mono);
|
||
}
|
||
Ok(resample_sinc(&mono, source_rate, SAMPLE_RATE))
|
||
}
|
||
|
||
fn resample_sinc(input: &[f32], source_rate: usize, target_rate: usize) -> Vec<f32> {
|
||
if input.len() < 2 || source_rate == 0 {
|
||
return input.to_vec();
|
||
}
|
||
let output_len = input
|
||
.len()
|
||
.saturating_mul(target_rate)
|
||
.checked_div(source_rate)
|
||
.unwrap_or(0)
|
||
.max(1);
|
||
let ratio = source_rate as f64 / target_rate as f64;
|
||
let cutoff = (target_rate as f64 / source_rate as f64).min(1.0) * 0.95;
|
||
const HALF_TAPS: isize = 8;
|
||
(0..output_len)
|
||
.map(|index| {
|
||
let position = index as f64 * ratio;
|
||
let center = position.floor() as isize;
|
||
let mut value = 0.0f64;
|
||
let mut weight_sum = 0.0f64;
|
||
for sample_index in center - HALF_TAPS + 1..=center + HALF_TAPS {
|
||
if sample_index < 0 || sample_index >= input.len() as isize {
|
||
continue;
|
||
}
|
||
let distance = position - sample_index as f64;
|
||
let phase = std::f64::consts::PI * distance * cutoff;
|
||
let sinc = if phase.abs() < 1e-12 {
|
||
1.0
|
||
} else {
|
||
phase.sin() / phase
|
||
};
|
||
let window_position = distance / HALF_TAPS as f64;
|
||
let window = if window_position.abs() <= 1.0 {
|
||
0.5 + 0.5 * (std::f64::consts::PI * window_position).cos()
|
||
} else {
|
||
0.0
|
||
};
|
||
let weight = cutoff * sinc * window;
|
||
value += input[sample_index as usize] as f64 * weight;
|
||
weight_sum += weight;
|
||
}
|
||
if weight_sum.abs() < 1e-12 {
|
||
input[center.clamp(0, input.len() as isize - 1) as usize]
|
||
} else {
|
||
(value / weight_sum) as f32
|
||
}
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
fn mel_spectrogram(signal: &[f32]) -> Result<Vec<[f32; MEL_BANDS]>> {
|
||
let frame_count = 1 + signal
|
||
.len()
|
||
.saturating_sub(FRAME_SIZE / 2)
|
||
.div_ceil(HOP_SIZE);
|
||
let filters = mel_filters();
|
||
let mut planner = FftPlanner::<f32>::new();
|
||
let fft = planner.plan_fft_forward(FRAME_SIZE);
|
||
let mut mel = Vec::with_capacity(frame_count);
|
||
let mut spectrum = vec![Complex::new(0.0f32, 0.0); FRAME_SIZE];
|
||
for frame_index in 0..frame_count {
|
||
let start = frame_index as isize * HOP_SIZE as isize - (FRAME_SIZE / 2) as isize;
|
||
for (index, value) in spectrum.iter_mut().enumerate() {
|
||
let source = start + index as isize;
|
||
let sample = if source >= 0 {
|
||
signal.get(source as usize).copied().unwrap_or(0.0)
|
||
} else {
|
||
0.0
|
||
};
|
||
let window = 0.5
|
||
- 0.5 * (2.0 * std::f32::consts::PI * index as f32 / (FRAME_SIZE - 1) as f32).cos();
|
||
*value = Complex::new(sample * window, 0.0);
|
||
}
|
||
fft.process(&mut spectrum);
|
||
let powers: Vec<f32> = spectrum[..=FRAME_SIZE / 2]
|
||
.iter()
|
||
.map(|value| value.norm_sqr())
|
||
.collect();
|
||
let mut bands = [0.0f32; MEL_BANDS];
|
||
for (band, weights) in filters.iter().enumerate() {
|
||
let energy: f32 = powers
|
||
.iter()
|
||
.zip(weights)
|
||
.map(|(power, weight)| power * weight)
|
||
.sum();
|
||
bands[band] = (1.0 + 10_000.0 * energy.max(0.0)).log10();
|
||
}
|
||
mel.push(bands);
|
||
}
|
||
Ok(mel)
|
||
}
|
||
|
||
fn mel_filters() -> Vec<Vec<f32>> {
|
||
let low = hz_to_mel_slaney(0.0);
|
||
let high = hz_to_mel_slaney((SAMPLE_RATE / 2) as f32);
|
||
let points: Vec<f32> = (0..MEL_BANDS + 2)
|
||
.map(|index| mel_to_hz_slaney(low + (high - low) * index as f32 / (MEL_BANDS + 1) as f32))
|
||
.collect();
|
||
let frequency_scale = (SAMPLE_RATE as f32 / 2.0) / (FRAME_SIZE / 2) as f32;
|
||
(0..MEL_BANDS)
|
||
.map(|band| {
|
||
let left = points[band];
|
||
let center = points[band + 1];
|
||
let right = points[band + 2];
|
||
let area = ((center - left) + (right - center)) / 2.0;
|
||
(0..=FRAME_SIZE / 2)
|
||
.map(|bin| {
|
||
let frequency = bin as f32 * frequency_scale;
|
||
let triangle = if frequency < left || frequency > right {
|
||
0.0
|
||
} else if frequency < center {
|
||
(frequency - left) / (center - left)
|
||
} else {
|
||
(right - frequency) / (right - center)
|
||
};
|
||
triangle.max(0.0) / area
|
||
})
|
||
.collect()
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
fn hz_to_mel_slaney(hz: f32) -> f32 {
|
||
if hz < 1000.0 {
|
||
hz / (200.0 / 3.0)
|
||
} else {
|
||
15.0 + 27.0 * (hz / 1000.0).ln() / 6.4f32.ln()
|
||
}
|
||
}
|
||
|
||
fn mel_to_hz_slaney(mel: f32) -> f32 {
|
||
if mel < 15.0 {
|
||
mel * (200.0 / 3.0)
|
||
} else {
|
||
1000.0 * (6.4f32.ln() * (mel - 15.0) / 27.0).exp()
|
||
}
|
||
}
|
||
|
||
fn normalize(vector: &mut [f32]) -> Result<()> {
|
||
let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
|
||
anyhow::ensure!(
|
||
norm.is_finite() && norm > f32::EPSILON,
|
||
"zero or invalid embedding"
|
||
);
|
||
for value in vector {
|
||
*value /= norm;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
fn dot(left: &[f32], right: &[f32]) -> f32 {
|
||
left.iter().zip(right).map(|(a, b)| a * b).sum()
|
||
}
|
||
|
||
fn is_near_duplicate(candidate: &[f32], kept: &[&[f32]]) -> bool {
|
||
kept.iter()
|
||
.any(|existing| dot(candidate, existing) >= NEAR_DUPLICATE_COSINE)
|
||
}
|
||
|
||
fn sha256_file(path: &Path) -> Result<String> {
|
||
use std::io::Read as _;
|
||
|
||
let mut file = File::open(path)?;
|
||
let mut hasher = Sha256::new();
|
||
let mut buffer = [0u8; 64 * 1024];
|
||
loop {
|
||
let read = file.read(&mut buffer)?;
|
||
if read == 0 {
|
||
break;
|
||
}
|
||
hasher.update(&buffer[..read]);
|
||
}
|
||
Ok(format!("{:x}", hasher.finalize()))
|
||
}
|
||
|
||
fn lock<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
|
||
mutex
|
||
.lock()
|
||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||
}
|
||
|
||
fn read<T>(lock: &RwLock<T>) -> std::sync::RwLockReadGuard<'_, T> {
|
||
lock.read()
|
||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||
}
|
||
|
||
fn write<T>(lock: &RwLock<T>) -> std::sync::RwLockWriteGuard<'_, T> {
|
||
lock.write()
|
||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
fn unique_test_dir(label: &str) -> PathBuf {
|
||
let unique = std::time::SystemTime::now()
|
||
.duration_since(std::time::UNIX_EPOCH)
|
||
.unwrap()
|
||
.as_nanos();
|
||
std::env::temp_dir().join(format!("furumi-{label}-{}-{unique}", std::process::id()))
|
||
}
|
||
|
||
#[test]
|
||
fn profile_fingerprint_changes_with_contract() {
|
||
let model = &MODELS[0];
|
||
let first = profile_fingerprint(model, DEFAULT_PROFILE_ID);
|
||
let second = profile_fingerprint(model, "another-profile");
|
||
assert_ne!(first, second);
|
||
assert_eq!(first, profile_fingerprint(model, DEFAULT_PROFILE_ID));
|
||
}
|
||
|
||
#[test]
|
||
fn profile_details_describe_the_processing_contract() {
|
||
let details = profile_details(DEFAULT_PROFILE_ID, DEFAULT_MODEL_ID).unwrap();
|
||
assert!(details.contains("Up to 300 seconds"));
|
||
assert!(details.contains("16000 Hz"));
|
||
assert!(details.contains("1280"));
|
||
}
|
||
|
||
#[test]
|
||
fn vectors_are_normalized() {
|
||
let mut vector = vec![3.0, 4.0];
|
||
normalize(&mut vector).unwrap();
|
||
assert!((dot(&vector, &vector) - 1.0).abs() < 1e-6);
|
||
}
|
||
|
||
#[test]
|
||
fn near_duplicate_embeddings_are_filtered_but_distinct_tracks_remain() {
|
||
let query = [1.0, 0.0, 0.0];
|
||
let near_duplicate = [0.99995, 0.01, 0.0];
|
||
let distinct = [0.0, 1.0, 0.0];
|
||
assert!(is_near_duplicate(&near_duplicate, &[&query]));
|
||
assert!(!is_near_duplicate(&distinct, &[&query]));
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn repeated_rescans_keep_an_up_to_date_profile_ready() {
|
||
let directory = unique_test_dir("similarity-stable-status");
|
||
let library = Arc::new(Library::open(&directory.join("library.db")).unwrap());
|
||
let profile_id = profile_fingerprint(&MODELS[0], DEFAULT_PROFILE_ID);
|
||
let settings = SimilaritySettings {
|
||
enabled: true,
|
||
active_profile: Some(profile_id),
|
||
..SimilaritySettings::default()
|
||
};
|
||
let (event_tx, mut event_rx) = tokio::sync::mpsc::unbounded_channel();
|
||
let manager = Manager::new(Arc::clone(&library), event_tx, settings);
|
||
|
||
assert_eq!(manager.status().phase, Phase::Ready);
|
||
for _ in 0..32 {
|
||
manager.start();
|
||
}
|
||
tokio::time::timeout(std::time::Duration::from_secs(2), async {
|
||
while manager.pipeline_running.load(Ordering::Acquire) {
|
||
tokio::task::yield_now().await;
|
||
}
|
||
})
|
||
.await
|
||
.unwrap();
|
||
|
||
assert_eq!(manager.status().phase, Phase::Ready);
|
||
while let Ok(event) = event_rx.try_recv() {
|
||
if let AppEvent::SimilarityStatus(status) = event {
|
||
assert_eq!(status.phase, Phase::Ready);
|
||
}
|
||
}
|
||
|
||
drop(manager);
|
||
drop(library);
|
||
std::fs::remove_dir_all(directory).unwrap();
|
||
}
|
||
|
||
#[test]
|
||
fn resampling_keeps_a_constant_signal() {
|
||
let output = resample_sinc(&vec![0.25; 441], 44_100, 16_000);
|
||
assert_eq!(output.len(), 160);
|
||
assert!(output.iter().all(|value| (*value - 0.25).abs() < 1e-6));
|
||
}
|
||
|
||
/// Manual compatibility check for the separately downloaded model:
|
||
/// `FURUMI_TEST_MODEL=/path/model.onnx cargo test onnx_model_smoke -- --ignored`.
|
||
#[test]
|
||
#[ignore]
|
||
fn onnx_model_smoke() {
|
||
let path = std::env::var_os("FURUMI_TEST_MODEL")
|
||
.map(PathBuf::from)
|
||
.expect("set FURUMI_TEST_MODEL");
|
||
let model = load_onnx(&path).unwrap();
|
||
let mut input = vec![0.0f32; MODEL_BATCH * PATCH_FRAMES * MEL_BANDS];
|
||
input[0] = 1.0;
|
||
let tensor = Tensor::from_shape(&[MODEL_BATCH, PATCH_FRAMES, MEL_BANDS], &input).unwrap();
|
||
let outputs = model.run(tvec!(tensor.clone().into_tvalue())).unwrap();
|
||
assert!(
|
||
outputs
|
||
.iter()
|
||
.any(|output| output.len() == MODEL_BATCH * EMBEDDING_DIMENSIONS)
|
||
);
|
||
let started = Instant::now();
|
||
for _ in 0..5 {
|
||
model.run(tvec!(tensor.clone().into_tvalue())).unwrap();
|
||
}
|
||
eprintln!(
|
||
"five warm batch-{MODEL_BATCH} runs: {:.3}s",
|
||
started.elapsed().as_secs_f64()
|
||
);
|
||
}
|
||
}
|