pub mod adapter;
pub mod download;
pub mod manifest;
pub mod metadata;
pub mod verify;
pub use adapter::{AdapterError, AdapterFactory, AdapterRegistry, AdapterStage, BuiltinAdapter};
pub use download::{
DownloadError, download_with_checksum, download_with_checksum_and_signature, verify_sha256,
};
pub use manifest::{
Manifest, ManifestError, ModelEntry, ProfileEntry, SCHEMA_V1, SCHEMA_V2, is_supported_schema,
};
pub use metadata::{MetaSource, ModelConfigMeta, load_model_config, read_onnx_metadata_props};
use crate::types::Profile;
use std::path::{Path, PathBuf};
pub const DEFAULT_MANIFEST_TOML: &str = include_str!("manifest.toml");
#[derive(Clone, Copy, Debug)]
pub struct VbxPldaArtifact {
pub id: &'static str,
pub filename: &'static str,
pub sha256: &'static str,
pub size: u64,
}
pub const VBX_PLDA_ARTIFACTS: &[VbxPldaArtifact] = &[
VbxPldaArtifact {
id: "vbx_plda_transform",
filename: "plda_transform.npy",
sha256: "90261469714415743f4b8a86ee6b89466db858bde3c5944367cccfb7abd34f14",
size: 131_200,
},
VbxPldaArtifact {
id: "vbx_plda_phi_computed",
filename: "plda_phi_computed.npy",
sha256: "6ef7cf2f5a23a45b66f440f9a996a4cf5c047b369829af695d50ef18aa0a35e3",
size: 1_152,
},
VbxPldaArtifact {
id: "vbx_plda_mean1",
filename: "plda_mean1.npy",
sha256: "e424c0c352182aa8e0f555dec1f3b30e29a20b9ed6b25d339f112af92e51e36f",
size: 2_176,
},
VbxPldaArtifact {
id: "vbx_plda_mean2",
filename: "plda_mean2.npy",
sha256: "6f6fb708a2037197b5b84ffeaa8f140cb878088fbecd6ab042ad26a7691bd2cf",
size: 640,
},
VbxPldaArtifact {
id: "vbx_plda_lda",
filename: "plda_lda.npy",
sha256: "e20c9b012bebd1aabda5a38a127e63a43cf35debdc502715fc143e2fb6bc3c4b",
size: 131_200,
},
VbxPldaArtifact {
id: "vbx_plda_mu",
filename: "plda_mu.npy",
sha256: "d286d48acf99bbc1ed1502fed0a3e361ae5626ce1870c8be9f7397c5e47886c6",
size: 1_152,
},
];
pub const VBX_PLDA_MODEL_IDS: &[&str] = &[
VBX_PLDA_ARTIFACTS[0].id,
VBX_PLDA_ARTIFACTS[1].id,
VBX_PLDA_ARTIFACTS[2].id,
VBX_PLDA_ARTIFACTS[3].id,
VBX_PLDA_ARTIFACTS[4].id,
VBX_PLDA_ARTIFACTS[5].id,
];
#[allow(clippy::expect_used)]
pub fn default_manifest() -> Manifest {
Manifest::from_toml_str(DEFAULT_MANIFEST_TOML)
.expect("embedded manifest.toml must parse — this is a static-asset bug")
}
#[derive(Debug, thiserror::Error)]
pub enum RegistryError {
#[error("model '{model_id}' not found in manifest")]
ModelNotFound { model_id: String },
#[error(
"model '{model_id}' has no signature in the manifest — release builds require a \
minisign signature for every profile-resolved model (a manifest that drops the \
signature would otherwise silently downgrade authenticity to a self-consistent hash)"
)]
UnsignedModel { model_id: String },
#[error("profile '{profile}' not found in manifest")]
ProfileNotFound { profile: String },
#[error("custom profile cannot be resolved by registry — caller must supply models")]
CustomProfileUnresolvable,
#[error("cache directory {path} is not writable")]
CacheNotWritable { path: PathBuf },
#[error("model '{model_id}' is not present in cache and offline mode is requested")]
OfflineMissing { model_id: String },
#[error("manifest error: {0}")]
Manifest(#[from] ManifestError),
#[error("download error: {0}")]
Download(#[from] DownloadError),
#[error("io error on {path}: {source}")]
Io {
path: PathBuf,
#[source]
source: std::io::Error,
},
}
#[derive(Debug, Clone)]
pub struct ProfileModels {
pub segmenter_path: PathBuf,
pub embedder_path: PathBuf,
}
#[derive(Debug, Clone)]
pub struct ModelRegistry {
manifest: Manifest,
cache_dir: PathBuf,
require_signatures: bool,
}
const REQUIRE_SIGNATURES_DEFAULT: bool = cfg!(not(debug_assertions));
impl ModelRegistry {
#[allow(clippy::should_implement_trait)]
pub fn default() -> Result<Self, RegistryError> {
let cache = dirs::cache_dir()
.ok_or_else(|| RegistryError::CacheNotWritable {
path: PathBuf::from("(unresolved-cache-dir)"),
})?
.join("polyvoice")
.join("models");
Self::with_cache_dir(cache)
}
pub fn with_cache_dir(path: impl AsRef<Path>) -> Result<Self, RegistryError> {
let path = path.as_ref().to_path_buf();
std::fs::create_dir_all(&path).map_err(|e| RegistryError::Io {
path: path.clone(),
source: e,
})?;
Ok(Self {
manifest: default_manifest(),
cache_dir: path,
require_signatures: REQUIRE_SIGNATURES_DEFAULT,
})
}
#[cfg(test)]
pub fn with_manifest_override(mut self, manifest: Manifest) -> Self {
self.manifest = manifest;
self
}
#[cfg(test)]
pub fn with_require_signatures(mut self, require: bool) -> Self {
self.require_signatures = require;
self
}
#[cfg(test)]
pub fn with_manifest(
manifest: Manifest,
cache_dir: impl AsRef<Path>,
) -> Result<Self, RegistryError> {
let path = cache_dir.as_ref().to_path_buf();
std::fs::create_dir_all(&path).map_err(|e| RegistryError::Io {
path: path.clone(),
source: e,
})?;
Ok(Self {
manifest,
cache_dir: path,
require_signatures: REQUIRE_SIGNATURES_DEFAULT,
})
}
pub fn cache_dir(&self) -> &Path {
&self.cache_dir
}
pub fn manifest(&self) -> &Manifest {
&self.manifest
}
pub fn ensure(&self, model_id: &str) -> Result<PathBuf, RegistryError> {
let entry = self
.manifest
.model(model_id)
.ok_or_else(|| RegistryError::ModelNotFound {
model_id: model_id.to_owned(),
})?;
let dest = self.cache_dir.join(&entry.filename);
download_with_checksum_and_signature(
&entry.url,
&entry.sha256,
entry.signature.as_deref(),
&dest,
)?;
Ok(dest)
}
pub fn ensure_vbx_plda_dir(&self) -> Result<PathBuf, RegistryError> {
for id in VBX_PLDA_MODEL_IDS {
self.ensure(id)?;
}
Ok(self.cache_dir.clone())
}
#[doc(hidden)]
#[cfg(test)] pub fn ensure_in_cache_only(&self, model_id: &str) -> Result<PathBuf, RegistryError> {
let entry = self
.manifest
.model(model_id)
.ok_or_else(|| RegistryError::ModelNotFound {
model_id: model_id.to_owned(),
})?;
let dest = self.cache_dir.join(&entry.filename);
if !dest.exists() {
return Err(RegistryError::OfflineMissing {
model_id: model_id.to_owned(),
});
}
Ok(dest)
}
fn require_signature_for(&self, model_id: &str) -> Result<(), RegistryError> {
if !self.require_signatures {
return Ok(());
}
let entry = self
.manifest
.model(model_id)
.ok_or_else(|| RegistryError::ModelNotFound {
model_id: model_id.to_owned(),
})?;
if entry.signature.is_none() {
return Err(RegistryError::UnsignedModel {
model_id: model_id.to_owned(),
});
}
Ok(())
}
pub fn ensure_for_profile(&self, profile: Profile) -> Result<ProfileModels, RegistryError> {
if profile == Profile::Custom {
return Err(RegistryError::CustomProfileUnresolvable);
}
let prof = self
.manifest
.profile(profile.manifest_id())
.ok_or_else(|| RegistryError::ProfileNotFound {
profile: profile.manifest_id().to_owned(),
})?;
self.require_signature_for(&prof.segmenter)?;
self.require_signature_for(&prof.embedder)?;
let segmenter_path = self.ensure(&prof.segmenter)?;
let embedder_path = self.ensure(&prof.embedder)?;
Ok(ProfileModels {
segmenter_path,
embedder_path,
})
}
#[cfg(test)]
pub fn ensure_in_cache_only_for_profile(
&self,
profile: Profile,
) -> Result<ProfileModels, RegistryError> {
if profile == Profile::Custom {
return Err(RegistryError::CustomProfileUnresolvable);
}
let prof = self
.manifest
.profile(profile.manifest_id())
.ok_or_else(|| RegistryError::ProfileNotFound {
profile: profile.manifest_id().to_owned(),
})?;
self.require_signature_for(&prof.segmenter)?;
self.require_signature_for(&prof.embedder)?;
let segmenter_path = self.ensure_in_cache_only(&prof.segmenter)?;
let embedder_path = self.ensure_in_cache_only(&prof.embedder)?;
Ok(ProfileModels {
segmenter_path,
embedder_path,
})
}
}
#[allow(clippy::unwrap_used)]
#[cfg(test)]
pub(crate) mod tests_helpers {
pub const TINY_MANIFEST: &str = r#"
schema = "polyvoice-models-v1"
[profiles.mobile]
segmenter = "hello_model"
embedder = "hello_model"
[profiles.balanced]
segmenter = "hello_model"
embedder = "hello_model"
[models.hello_model]
url = "file:///dev/null"
sha256 = "2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824"
size = 5
filename = "hello.bin"
"#;
}
#[allow(clippy::unwrap_used)]
#[cfg(test)]
mod tests {
use super::*;
use crate::types::Profile;
use std::path::Path;
use tempfile::TempDir;
#[test]
fn embedded_manifest_parses() {
let m = default_manifest();
assert!(
is_supported_schema(&m.schema),
"embedded schema must be v1 or v2, got {}",
m.schema
);
assert!(m.profiles.contains_key("mobile"));
assert!(m.profiles.contains_key("balanced"));
}
#[test]
fn embedded_manifest_is_v2_with_adapter_metadata() {
let m = default_manifest();
assert_eq!(m.schema, SCHEMA_V2);
for (id, entry) in &m.models {
assert!(
entry.adapter_type.is_some(),
"model '{id}' missing adapter_type"
);
assert!(entry.license.is_some(), "model '{id}' missing license");
assert!(entry.version.is_some(), "model '{id}' missing version");
}
assert_eq!(
m.resolve_model_ref("segmenter", "latest"),
Some("powerset_fp32")
);
assert_eq!(
m.resolve_model_ref("embedder", "latest"),
Some("wespeaker_resnet34")
);
assert_eq!(m.resolve_model_ref("vad", "latest"), Some("silero_vad"));
}
#[test]
fn embedded_manifest_lists_legacy_models() {
let m = default_manifest();
assert!(m.models.contains_key("silero_vad"));
assert!(m.models.contains_key("wespeaker_resnet34"));
}
#[test]
fn profiles_share_segmenter_and_embedder_in_v2_hotfix() {
let m = default_manifest();
let mob = m.profile("mobile").unwrap();
let bal = m.profile("balanced").unwrap();
assert_eq!(mob.segmenter, bal.segmenter, "both use powerset");
assert_eq!(
mob.embedder, bal.embedder,
"both use resnet34 (CAM++ broken)"
);
}
#[test]
fn registry_default_uses_user_cache() {
let r = ModelRegistry::default().expect("default cache dir resolvable");
let path = r.cache_dir().to_path_buf();
assert!(path.ends_with("polyvoice/models") || path.ends_with("polyvoice\\models"));
}
#[test]
fn registry_with_cache_dir_creates_dir() {
let tmp = TempDir::new().unwrap();
let path = tmp.path().join("nested/models");
let r = ModelRegistry::with_cache_dir(&path).unwrap();
assert!(path.exists());
assert_eq!(r.cache_dir(), path.as_path());
}
#[test]
fn ensure_returns_err_for_unknown_model_id() {
let tmp = TempDir::new().unwrap();
let r = ModelRegistry::with_cache_dir(tmp.path()).unwrap();
let err = r
.ensure_in_cache_only("ghost")
.expect_err("must be missing");
assert!(matches!(err, RegistryError::ModelNotFound { .. }));
}
#[test]
fn ensure_in_cache_only_succeeds_when_file_present() {
let tmp = TempDir::new().unwrap();
let manifest =
Manifest::from_toml_str(crate::models::tests_helpers::TINY_MANIFEST).unwrap();
let r = ModelRegistry::with_cache_dir(tmp.path())
.unwrap()
.with_manifest_override(manifest);
let cached = tmp.path().join("hello.bin");
std::fs::write(&cached, b"hello").unwrap();
let path = r.ensure_in_cache_only("hello_model").unwrap();
assert_eq!(path, cached);
}
#[test]
fn ensure_for_profile_uses_manifest_lookup() {
let tmp = TempDir::new().unwrap();
let manifest =
Manifest::from_toml_str(crate::models::tests_helpers::TINY_MANIFEST).unwrap();
let r = ModelRegistry::with_cache_dir(tmp.path())
.unwrap()
.with_manifest_override(manifest)
.with_require_signatures(false);
std::fs::write(tmp.path().join("hello.bin"), b"hello").unwrap();
let bundle = r.ensure_in_cache_only_for_profile(Profile::Mobile).unwrap();
assert_eq!(bundle.segmenter_path, tmp.path().join("hello.bin"));
assert_eq!(bundle.embedder_path, tmp.path().join("hello.bin"));
}
const TINY_MANIFEST_SIGNED: &str = r#"
schema = "polyvoice-models-v1"
[profiles.mobile]
segmenter = "hello_model"
embedder = "hello_model"
[profiles.balanced]
segmenter = "hello_model"
embedder = "hello_model"
[models.hello_model]
url = "file:///dev/null"
sha256 = "2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824"
size = 5
filename = "hello.bin"
signature = "untrusted comment: fixture\nRWQfixturesignature"
"#;
#[test]
fn strict_profile_resolution_rejects_unsigned_model() {
let tmp = TempDir::new().unwrap();
let manifest =
Manifest::from_toml_str(crate::models::tests_helpers::TINY_MANIFEST).unwrap();
let r = ModelRegistry::with_cache_dir(tmp.path())
.unwrap()
.with_manifest_override(manifest)
.with_require_signatures(true);
let err = r.ensure_for_profile(Profile::Mobile).expect_err("unsigned");
assert!(
matches!(err, RegistryError::UnsignedModel { ref model_id } if model_id == "hello_model")
);
let err = r
.ensure_in_cache_only_for_profile(Profile::Mobile)
.expect_err("unsigned");
assert!(matches!(err, RegistryError::UnsignedModel { .. }));
}
#[test]
fn strict_profile_resolution_accepts_signed_model() {
let tmp = TempDir::new().unwrap();
let manifest = Manifest::from_toml_str(TINY_MANIFEST_SIGNED).unwrap();
let r = ModelRegistry::with_cache_dir(tmp.path())
.unwrap()
.with_manifest_override(manifest)
.with_require_signatures(true);
std::fs::write(tmp.path().join("hello.bin"), b"hello").unwrap();
let bundle = r.ensure_in_cache_only_for_profile(Profile::Mobile).unwrap();
assert_eq!(bundle.segmenter_path, tmp.path().join("hello.bin"));
}
#[test]
fn every_profile_model_is_signed() {
let m = default_manifest();
for (profile_id, prof) in &m.profiles {
for model_id in [&prof.segmenter, &prof.embedder] {
let entry = m.models.get(model_id).unwrap_or_else(|| {
panic!("profile '{profile_id}' references missing model '{model_id}'")
});
assert!(
entry.signature.is_some(),
"profile model '{model_id}' (via '{profile_id}') has no signature — \
release profile resolution would fail"
);
}
}
}
#[test]
fn optional_sortformer_entry_present_but_not_in_profiles() {
let m = default_manifest();
let entry = m.model("sortformer_v2").expect("sortformer_v2 in manifest");
assert_eq!(entry.adapter_type.as_deref(), Some("sortformer-v2"));
assert_eq!(entry.license.as_deref(), Some("CC-BY-4.0"));
assert_eq!(entry.num_speakers, Some(4));
for (pid, prof) in &m.profiles {
assert_ne!(
prof.segmenter, "sortformer_v2",
"profile {pid} must not pull sortformer as segmenter"
);
assert_ne!(
prof.embedder, "sortformer_v2",
"profile {pid} must not pull sortformer as embedder"
);
}
}
#[test]
fn optional_vbx_plda_entries_present_but_not_in_profiles() {
let m = default_manifest();
assert_eq!(VBX_PLDA_MODEL_IDS.len(), VBX_PLDA_ARTIFACTS.len());
for (art, listed_id) in VBX_PLDA_ARTIFACTS.iter().zip(VBX_PLDA_MODEL_IDS.iter()) {
assert_eq!(art.id, *listed_id);
let entry = m
.model(art.id)
.unwrap_or_else(|| panic!("missing manifest entry {}", art.id));
assert_eq!(entry.sha256, art.sha256, "{} sha256 mismatch", art.id);
assert_eq!(entry.size, Some(art.size), "{} size mismatch", art.id);
assert_eq!(entry.filename, art.filename, "{} filename mismatch", art.id);
assert_eq!(entry.adapter_type.as_deref(), Some("vbx-plda"));
assert_eq!(entry.license.as_deref(), Some("CC-BY-4.0"));
assert!(
entry.signature.is_none(),
"{} must stay unsigned until a release engineer signs it \
(profile resolution never pulls these)",
art.id
);
assert!(
entry.url.starts_with("https://"),
"{} url must be https",
art.id
);
for (pid, prof) in &m.profiles {
assert_ne!(
prof.segmenter.as_str(),
art.id,
"profile {pid} must not pull PLDA as segmenter"
);
assert_ne!(
prof.embedder.as_str(),
art.id,
"profile {pid} must not pull PLDA as embedder"
);
}
}
}
#[test]
fn ensure_vbx_plda_dir_uses_local_cache_without_network() {
let tmp = TempDir::new().unwrap();
let fixture_dir = Path::new(env!("CARGO_MANIFEST_DIR")).join("fixtures/vbx-plda");
for art in VBX_PLDA_ARTIFACTS {
let src = fixture_dir.join(art.filename);
assert!(
src.is_file(),
"fixture missing: {} (run scripts/build-vbx-plda.py)",
src.display()
);
std::fs::copy(&src, tmp.path().join(art.filename)).unwrap();
}
let r = ModelRegistry::with_cache_dir(tmp.path()).unwrap();
let dir = r
.ensure_vbx_plda_dir()
.expect("cache-hit ensure must succeed offline");
assert_eq!(dir, tmp.path());
for art in VBX_PLDA_ARTIFACTS {
assert!(dir.join(art.filename).is_file());
}
}
}