speakrs 0.7.0

use std::path::{Path, PathBuf};

use crate::inference::ExecutionMode;
#[cfg(feature = "online")]
use hf_hub::api::sync::{Api, ApiBuilder, ApiRepo};
#[cfg(feature = "online")]
use hf_hub::{Repo, RepoType};

const SEGMENTATION_ONNX: &str = "segmentation-3.0.onnx";
const EMBEDDING_ONNX: &str = "wespeaker-voxceleb-resnet34.onnx";
const EMBEDDING_MIN_SAMPLES: &str = "wespeaker-voxceleb-resnet34.min_num_samples.txt";

/// Resolved model paths for the speakrs pipeline
///
/// Captures the three root paths needed by [`SegmentationModel`], [`EmbeddingModel`],
/// and `PldaTransform`. Native CPU, CUDA and CoreML models use the base ONNX paths
/// as family selectors; the ONNX files do not need to exist for those modes. CPU and
/// CUDA load `segmentation-3.0.safetensors` and `wespeaker-multimask-tail.safetensors`
/// from the same directory. MIGraphX loads the ONNX files
///
/// [`SegmentationModel`]: crate::inference::segmentation::SegmentationModel
/// [`EmbeddingModel`]: crate::inference::embedding::EmbeddingModel
#[derive(Debug, Clone)]
pub struct ModelBundle {
    segmentation_onnx: PathBuf,
    embedding_onnx: PathBuf,
    plda_dir: PathBuf,
}

impl ModelBundle {
    /// Resolve and validate paths from a local directory containing all model files
    pub fn from_dir(
        models_dir: impl Into<PathBuf>,
    ) -> Result<Self, crate::inference::ModelLoadError> {
        let dir = models_dir.into();
        crate::inference::embedding::read_min_num_samples(&dir.join(EMBEDDING_MIN_SAMPLES))?;

        Ok(Self {
            segmentation_onnx: dir.join(SEGMENTATION_ONNX),
            embedding_onnx: dir.join(EMBEDDING_ONNX),
            plda_dir: dir,
        })
    }

    /// Download models from HuggingFace and resolve paths
    #[cfg(feature = "online")]
    #[cfg_attr(docsrs, doc(cfg(feature = "online")))]
    pub fn from_pretrained(mode: ExecutionMode) -> Result<Self, crate::inference::ModelLoadError> {
        let manager = ModelManager::new()?;
        let dir = manager
            .ensure(mode)
            .map_err(|source| download_error(mode, source))?;
        Self::from_dir(dir)
    }

    /// Segmentation family selector, or the ONNX model path for MIGraphX
    pub fn segmentation_path(&self) -> &Path {
        &self.segmentation_onnx
    }

    /// Embedding family selector, or the ONNX model path for MIGraphX
    pub fn embedding_path(&self) -> &Path {
        &self.embedding_onnx
    }

    /// Directory containing PLDA parameter files
    pub fn plda_dir(&self) -> &Path {
        &self.plda_dir
    }
}

/// Names the CUDA weights in the error when they cannot be downloaded
#[cfg(feature = "online")]
fn download_error(
    mode: ExecutionMode,
    source: hf_hub::api::sync::ApiError,
) -> crate::inference::ModelLoadError {
    #[cfg(feature = "_cuda")]
    if mode.is_cuda() {
        return crate::inference::ModelLoadError::CudaAssetsUnavailable { mode, source };
    }

    let _ = mode;
    source.into()
}

/// Hugging Face repository that hosts the speakrs models
pub const MODEL_REPOSITORY: &str = "avencera/speakrs-models";
/// Revision of [`MODEL_REPOSITORY`] that this release loads
// CI downloads fixtures from this same revision via SPEAKRS_MODEL_FIXTURE_REV
pub const MODEL_REVISION: &str = "67f963b6a1d94fcd4678948237b3cc8bd69bf1d8";

#[cfg(feature = "online")]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct ModelRepositoryId(&'static str);

#[cfg(feature = "online")]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct ModelRevision(&'static str);

#[cfg(feature = "online")]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct PinnedModelRepository {
    id: ModelRepositoryId,
    revision: ModelRevision,
}

#[cfg(feature = "online")]
impl PinnedModelRepository {
    const fn speakrs_models() -> Self {
        Self {
            id: ModelRepositoryId(MODEL_REPOSITORY),
            revision: ModelRevision(MODEL_REVISION),
        }
    }

    fn into_hf_repo(self) -> Repo {
        Repo::with_revision(
            self.id.0.to_owned(),
            RepoType::Model,
            self.revision.0.to_owned(),
        )
    }

    fn bind(self, api: &Api) -> ApiRepo {
        api.repo(self.into_hf_repo())
    }
}

#[cfg(feature = "online")]
const PINNED_MODEL_REPOSITORY: PinnedModelRepository = PinnedModelRepository::speakrs_models();

/// Manages downloading and caching speakrs models from HuggingFace
#[cfg(feature = "online")]
#[cfg_attr(docsrs, doc(cfg(feature = "online")))]
pub struct ModelManager {
    repo: ApiRepo,
}

#[cfg(feature = "online")]
impl ModelManager {
    /// Create a manager using the default HuggingFace cache directory
    pub fn new() -> Result<Self, hf_hub::api::sync::ApiError> {
        let api = Api::new()?;
        Ok(Self::from_api(api))
    }

    /// Create a manager with a custom cache directory
    pub fn with_cache_dir(cache_dir: PathBuf) -> Result<Self, hf_hub::api::sync::ApiError> {
        let api = ApiBuilder::from_cache(hf_hub::Cache::new(cache_dir)).build()?;
        Ok(Self::from_api(api))
    }

    fn from_api(api: Api) -> Self {
        let repo = PINNED_MODEL_REPOSITORY.bind(&api);
        Self { repo }
    }

    /// Download a single file, returns path to cached copy
    pub fn get(&self, filename: impl AsRef<str>) -> Result<PathBuf, hf_hub::api::sync::ApiError> {
        self.repo.get(filename.as_ref())
    }

    /// Ensure all files for a mode are downloaded, return base models dir
    pub fn ensure(&self, mode: ExecutionMode) -> Result<PathBuf, hf_hub::api::sync::ApiError> {
        let files = required_model_files(mode);
        for file in &files {
            self.repo.get(file)?;
        }
        // all files land in the same snapshot dir
        let first = self.repo.get(&files[0])?;
        let Some(parent) = first.parent() else {
            return Ok(first);
        };
        Ok(parent.to_path_buf())
    }
}

const PLDA_FILES: &[&str] = &[
    "plda_lda.npy",
    "plda_tr.npy",
    "plda_mu.npy",
    "plda_psi.npy",
    "plda_mean1.npy",
    "plda_mean2.npy",
    EMBEDDING_MIN_SAMPLES,
];

/// The safetensors weights the native CPU and CUDA modes load, exported by
/// `scripts/cuda/export_weights.py --runtime-assets`
const NATIVE_WEIGHT_FILES: &[(&str, ModelFamily)] = &[
    ("segmentation-3.0.safetensors", ModelFamily::Segmentation),
    (
        "wespeaker-multimask-tail.safetensors",
        ModelFamily::MultiMaskTail,
    ),
];

fn mlmodelc_files(name: &str) -> Vec<String> {
    vec![
        format!("{name}/model.mil"),
        format!("{name}/coremldata.bin"),
        format!("{name}/weights/weight.bin"),
        format!("{name}/analytics/coremldata.bin"),
    ]
}

const COREML_COMMON_MODEL_STEMS: &[&str] = &[
    "segmentation-3.0.mlmodelc",
    "segmentation-3.0-b32.mlmodelc",
    "segmentation-3.0-b64.mlmodelc",
    "wespeaker-fbank.mlmodelc",
    "wespeaker-fbank-b32.mlmodelc",
    "wespeaker-fbank-30s.mlmodelc",
    "wespeaker-multimask-tail-b32.mlmodelc",
    "wespeaker-voxceleb-resnet34-tail.mlmodelc",
    "wespeaker-voxceleb-resnet34-tail-b3.mlmodelc",
    "wespeaker-voxceleb-resnet34-tail-b32.mlmodelc",
];

const COREML_CHUNK_MODEL_STEMS: &[&str] = &[
    "wespeaker-chunk-emb-p1s-w21.mlmodelc",
    "wespeaker-chunk-emb-p1s-w36.mlmodelc",
    "wespeaker-chunk-emb-p1s-w51.mlmodelc",
    "wespeaker-chunk-emb-p1s-w81.mlmodelc",
    "wespeaker-chunk-emb-p1s-w111.mlmodelc",
];

const COREML_FAST_SEGMENTATION_MODEL_STEMS: &[&str] = &[
    "segmentation-3.0-w8a16.mlmodelc",
    "segmentation-3.0-b32-w8a16.mlmodelc",
    "segmentation-3.0-b64-w8a16.mlmodelc",
];

const COREML_FAST_CHUNK_MODEL_STEMS: &[&str] = &[
    "wespeaker-chunk-emb-s25-w11.mlmodelc",
    "wespeaker-chunk-emb-s25-w16.mlmodelc",
    "wespeaker-chunk-emb-s25-w21.mlmodelc",
    "wespeaker-chunk-emb-s25-w26.mlmodelc",
    "wespeaker-chunk-emb-s25-w36.mlmodelc",
    "wespeaker-chunk-emb-s25-w46.mlmodelc",
    "wespeaker-chunk-emb-s25-w56.mlmodelc",
];

/// Model family owned by the asset catalog
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum ModelFamily {
    Segmentation,
    Embedding,
    Filterbank,
    EmbeddingTail,
    MultiMaskTail,
    ChunkEmbedding,
    Plda,
    Metadata,
}

/// Backend that consumes a catalog asset
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum ModelBackend {
    SharedData,
    Cpu,
    Onnx,
    CoreMl,
    Cuda,
}

/// Weight precision recorded for a catalog asset
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum ModelPrecision {
    Fp32,
    W8A16,
}

/// Whether a catalog asset must be downloaded for a mode
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum AssetRequirement {
    Required,
    Optional,
}

/// One remote model file or compiled CoreML bundle stem
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct ModelAsset {
    file_name: &'static str,
    family: ModelFamily,
    backend: ModelBackend,
    precision: ModelPrecision,
    batch: Option<u32>,
    requirement: AssetRequirement,
}

impl ModelAsset {
    const fn new(
        file_name: &'static str,
        family: ModelFamily,
        backend: ModelBackend,
        precision: ModelPrecision,
        batch: Option<u32>,
        requirement: AssetRequirement,
    ) -> Self {
        Self {
            file_name,
            family,
            backend,
            precision,
            batch,
            requirement,
        }
    }

    /// Remote file name or compiled-bundle directory
    #[cfg_attr(not(test), allow(dead_code))]
    pub(crate) fn file_name(self) -> &'static str {
        self.file_name
    }

    /// Model family this asset belongs to
    #[cfg_attr(not(test), allow(dead_code))]
    pub(crate) fn family(self) -> ModelFamily {
        self.family
    }

    /// Backend that loads this asset
    #[cfg_attr(not(test), allow(dead_code))]
    pub(crate) fn backend(self) -> ModelBackend {
        self.backend
    }

    /// Recorded weight precision
    #[cfg_attr(not(test), allow(dead_code))]
    pub(crate) fn precision(self) -> ModelPrecision {
        self.precision
    }

    /// Batch dimension when the asset is a fixed-shape model
    #[cfg_attr(not(test), allow(dead_code))]
    pub(crate) fn batch(self) -> Option<u32> {
        self.batch
    }

    /// Whether this asset must be present for the selected mode
    pub(crate) fn requirement(self) -> AssetRequirement {
        self.requirement
    }

    fn remote_paths(self) -> Vec<String> {
        if self.file_name.ends_with(".mlmodelc") {
            mlmodelc_files(self.file_name)
        } else {
            vec![self.file_name.to_string()]
        }
    }
}

fn family_for_stem(name: &str) -> ModelFamily {
    if name.starts_with("segmentation") {
        ModelFamily::Segmentation
    } else if name.contains("chunk-emb") {
        ModelFamily::ChunkEmbedding
    } else if name.contains("fbank") {
        ModelFamily::Filterbank
    } else if name.contains("multimask") {
        ModelFamily::MultiMaskTail
    } else if name.contains("tail") {
        ModelFamily::EmbeddingTail
    } else {
        ModelFamily::Embedding
    }
}

fn batch_from_name(name: &str) -> Option<u32> {
    let stem = name
        .trim_end_matches(".mlmodelc")
        .trim_end_matches(".onnx")
        .trim_end_matches(".onnx.data");
    stem.rsplit_once("-b")
        .and_then(|(_, batch)| batch.parse().ok())
}

fn onnx_asset(name: &'static str, family: ModelFamily) -> ModelAsset {
    ModelAsset::new(
        name,
        family,
        ModelBackend::Onnx,
        ModelPrecision::Fp32,
        batch_from_name(name),
        AssetRequirement::Required,
    )
}

fn catalog_assets(mode: ExecutionMode) -> Vec<ModelAsset> {
    let mut assets: Vec<ModelAsset> = PLDA_FILES
        .iter()
        .map(|name| {
            let family = if name.ends_with(".txt") {
                ModelFamily::Metadata
            } else {
                ModelFamily::Plda
            };
            ModelAsset::new(
                name,
                family,
                ModelBackend::SharedData,
                ModelPrecision::Fp32,
                None,
                AssetRequirement::Required,
            )
        })
        .collect();

    match mode {
        // native modes share the deployed FP32 weights and never read ONNX graphs
        ExecutionMode::Cpu | ExecutionMode::Cuda | ExecutionMode::CudaFast => {
            let backend = if matches!(mode, ExecutionMode::Cpu) {
                ModelBackend::Cpu
            } else {
                ModelBackend::Cuda
            };
            assets.extend(NATIVE_WEIGHT_FILES.iter().map(|&(name, family)| {
                ModelAsset::new(
                    name,
                    family,
                    backend,
                    ModelPrecision::Fp32,
                    None,
                    AssetRequirement::Required,
                )
            }));
        }
        ExecutionMode::MiGraphX => {
            assets.push(onnx_asset(
                "segmentation-3.0.onnx",
                ModelFamily::Segmentation,
            ));
            assets.push(onnx_asset(
                "wespeaker-voxceleb-resnet34.onnx",
                ModelFamily::Embedding,
            ));
            assets.push(onnx_asset(
                "wespeaker-voxceleb-resnet34.onnx.data",
                ModelFamily::Embedding,
            ));
            assets.push(onnx_asset("wespeaker-fbank.onnx", ModelFamily::Filterbank));
            assets.push(onnx_asset(
                "wespeaker-fbank-b32.onnx",
                ModelFamily::Filterbank,
            ));
            assets.push(onnx_asset(
                "wespeaker-multimask-tail.onnx",
                ModelFamily::MultiMaskTail,
            ));
            assets.push(onnx_asset(
                "wespeaker-multimask-tail-b32.onnx",
                ModelFamily::MultiMaskTail,
            ));
            assets.push(onnx_asset(
                "segmentation-3.0-b32.onnx",
                ModelFamily::Segmentation,
            ));
            assets.push(onnx_asset(
                "wespeaker-voxceleb-resnet34-b64.onnx",
                ModelFamily::Embedding,
            ));
        }
        // compiled bundles only: CoreML modes never read the ONNX files
        ExecutionMode::CoreMl | ExecutionMode::CoreMlFast => {
            assets.extend(COREML_COMMON_MODEL_STEMS.iter().map(|&name| {
                ModelAsset::new(
                    name,
                    family_for_stem(name),
                    ModelBackend::CoreMl,
                    ModelPrecision::Fp32,
                    batch_from_name(name),
                    AssetRequirement::Required,
                )
            }));
            if matches!(mode, ExecutionMode::CoreMl) {
                assets.extend(COREML_CHUNK_MODEL_STEMS.iter().map(|&name| {
                    ModelAsset::new(
                        name,
                        ModelFamily::ChunkEmbedding,
                        ModelBackend::CoreMl,
                        ModelPrecision::Fp32,
                        batch_from_name(name),
                        AssetRequirement::Required,
                    )
                }));
            }
            if matches!(mode, ExecutionMode::CoreMlFast) {
                assets.extend(COREML_FAST_SEGMENTATION_MODEL_STEMS.iter().map(|&name| {
                    ModelAsset::new(
                        name,
                        ModelFamily::Segmentation,
                        ModelBackend::CoreMl,
                        ModelPrecision::W8A16,
                        batch_from_name(name),
                        AssetRequirement::Required,
                    )
                }));
                assets.extend(COREML_FAST_CHUNK_MODEL_STEMS.iter().map(|&name| {
                    ModelAsset::new(
                        name,
                        ModelFamily::ChunkEmbedding,
                        ModelBackend::CoreMl,
                        ModelPrecision::Fp32,
                        batch_from_name(name),
                        AssetRequirement::Required,
                    )
                }));
            }
            assets.push(ModelAsset::new(
                "wespeaker-voxceleb-resnet34-tail-b64.mlmodelc",
                ModelFamily::EmbeddingTail,
                ModelBackend::CoreMl,
                ModelPrecision::Fp32,
                Some(64),
                AssetRequirement::Optional,
            ));
        }
    }

    assets
}

/// Files that `mode` loads, relative to the model directory
///
/// The model manager downloads exactly these files from [`MODEL_REPOSITORY`] at
/// [`MODEL_REVISION`]. Compiled CoreML bundles are listed file by file
pub fn required_model_files(mode: ExecutionMode) -> Vec<String> {
    catalog_assets(mode)
        .into_iter()
        .filter(|asset| asset.requirement() == AssetRequirement::Required)
        .flat_map(ModelAsset::remote_paths)
        .collect()
}

#[cfg(test)]
mod local_tests {
    use std::fs;
    use std::sync::atomic::{AtomicU64, Ordering};

    use super::{EMBEDDING_MIN_SAMPLES, ModelBundle};

    static NEXT_DIRECTORY_ID: AtomicU64 = AtomicU64::new(0);

    fn scratch_directory(test_name: &str) -> std::path::PathBuf {
        let id = NEXT_DIRECTORY_ID.fetch_add(1, Ordering::Relaxed);
        let directory = std::env::temp_dir().join(format!(
            "speakrs-model-bundle-{}-{test_name}-{id}",
            std::process::id()
        ));
        fs::create_dir_all(&directory).unwrap();
        directory
    }

    #[test]
    fn local_bundle_requires_embedding_metadata() {
        let directory = scratch_directory("missing-metadata");
        let error = ModelBundle::from_dir(&directory).unwrap_err();

        assert!(error.to_string().contains("missing embedding metadata"));
        fs::remove_dir_all(directory).unwrap();
    }

    #[test]
    fn local_bundle_validates_embedding_metadata_before_construction() {
        let directory = scratch_directory("metadata-validation");
        let metadata = directory.join(EMBEDDING_MIN_SAMPLES);
        fs::write(&metadata, "0\n").unwrap();
        let error = ModelBundle::from_dir(&directory).unwrap_err();
        assert!(error.to_string().contains("must be greater than zero"));

        fs::write(metadata, "400\n").unwrap();
        ModelBundle::from_dir(&directory).unwrap();
        fs::remove_dir_all(directory).unwrap();
    }
}

// the catalog tests also keep its accessors used in builds without `online`
#[cfg(test)]
mod tests {
    use super::*;

    #[cfg(feature = "online")]
    const MODEL_FILENAME: &str = "segmentation-3.0.onnx";
    #[cfg(feature = "online")]
    const EXPECTED_MODEL_URL: &str = "https://huggingface.co/avencera/speakrs-models/resolve/67f963b6a1d94fcd4678948237b3cc8bd69bf1d8/segmentation-3.0.onnx";

    #[cfg(feature = "online")]
    #[test]
    fn pinned_repository_selects_model_card_revision() {
        let repo = PINNED_MODEL_REPOSITORY.into_hf_repo();
        assert_eq!(repo.folder_name(), "models--avencera--speakrs-models");
        assert_eq!(repo.revision(), MODEL_REVISION);

        let api = ApiBuilder::new().with_progress(false).build().unwrap();
        let api_repo = api.repo(repo);
        assert_eq!(api_repo.url(MODEL_FILENAME), EXPECTED_MODEL_URL);
    }

    #[cfg(feature = "online")]
    #[test]
    fn default_and_custom_cache_managers_share_pinned_repository() {
        let default_manager = ModelManager::new().unwrap();
        let custom_cache = std::env::temp_dir().join(format!(
            "speakrs-pinned-model-manager-{}",
            std::process::id()
        ));
        let custom_manager = ModelManager::with_cache_dir(custom_cache).unwrap();

        assert_eq!(default_manager.repo.url(MODEL_FILENAME), EXPECTED_MODEL_URL);
        assert_eq!(custom_manager.repo.url(MODEL_FILENAME), EXPECTED_MODEL_URL);
    }

    #[test]
    fn coreml_required_files_include_chunk_fast_path_assets() {
        let files = required_model_files(ExecutionMode::CoreMl);
        assert!(files.contains(&"segmentation-3.0-b64.mlmodelc/model.mil".to_string()));
        assert!(files.contains(&"wespeaker-fbank-30s.mlmodelc/model.mil".to_string()));
        assert!(files.contains(&"wespeaker-multimask-tail-b32.mlmodelc/model.mil".to_string()));
        assert!(files.contains(&"wespeaker-chunk-emb-p1s-w111.mlmodelc/model.mil".to_string()));
    }

    #[test]
    fn coreml_modes_download_no_onnx_files() {
        for mode in [ExecutionMode::CoreMl, ExecutionMode::CoreMlFast] {
            let files = required_model_files(mode);
            assert!(
                files.iter().all(|path| !path.contains(".onnx")),
                "{mode} downloads ONNX files: {files:?}"
            );
            assert!(files.contains(&EMBEDDING_MIN_SAMPLES.to_string()));
            assert!(files.contains(&"wespeaker-fbank.mlmodelc/model.mil".to_string()));
        }
    }

    #[test]
    fn coreml_fast_required_files_include_fast_assets() {
        let files = required_model_files(ExecutionMode::CoreMlFast);
        assert!(files.contains(&"segmentation-3.0-w8a16.mlmodelc/model.mil".to_string()));
        assert!(files.contains(&"segmentation-3.0-b64-w8a16.mlmodelc/model.mil".to_string()));
        assert!(files.contains(&"wespeaker-chunk-emb-s25-w56.mlmodelc/model.mil".to_string()));
    }

    #[test]
    fn catalog_selects_required_batch_32_coreml_tail() {
        let tail = catalog_assets(ExecutionMode::CoreMl)
            .into_iter()
            .find(|asset| asset.family() == ModelFamily::EmbeddingTail && asset.batch() == Some(32))
            .expect("required batch-32 CoreML tail");
        assert_eq!(tail.backend(), ModelBackend::CoreMl);
        assert_eq!(tail.precision(), ModelPrecision::Fp32);
        assert_eq!(tail.requirement(), AssetRequirement::Required);
        assert_eq!(
            tail.file_name(),
            "wespeaker-voxceleb-resnet34-tail-b32.mlmodelc"
        );
        assert_eq!(
            tail.remote_paths()[0],
            "wespeaker-voxceleb-resnet34-tail-b32.mlmodelc/model.mil"
        );
    }

    #[test]
    fn coreml_catalog_records_optional_batch_64_tail() {
        let optional = catalog_assets(ExecutionMode::CoreMl)
            .into_iter()
            .find(|asset| asset.batch() == Some(64) && asset.family() == ModelFamily::EmbeddingTail)
            .expect("optional batch-64 CoreML tail");
        assert_eq!(optional.requirement(), AssetRequirement::Optional);
        let files = required_model_files(ExecutionMode::CoreMl);
        assert!(!files.iter().any(|path| path.contains("tail-b64")));
    }

    #[test]
    fn native_modes_download_weights_plda_and_metadata_without_onnx() {
        let mut expected: Vec<String> = PLDA_FILES
            .iter()
            .map(|name| name.to_string())
            .chain([
                "segmentation-3.0.safetensors".to_owned(),
                "wespeaker-multimask-tail.safetensors".to_owned(),
            ])
            .collect();
        expected.sort();

        for mode in [
            ExecutionMode::Cpu,
            ExecutionMode::Cuda,
            ExecutionMode::CudaFast,
        ] {
            let mut files = required_model_files(mode);
            files.sort();
            assert_eq!(files, expected, "unexpected assets for {mode}");

            let backend = if matches!(mode, ExecutionMode::Cpu) {
                ModelBackend::Cpu
            } else {
                ModelBackend::Cuda
            };
            let assets = catalog_assets(mode);
            assert!(assets.iter().any(|asset| {
                asset.family() == ModelFamily::Segmentation && asset.backend() == backend
            }));
            assert!(assets.iter().any(|asset| {
                asset.family() == ModelFamily::MultiMaskTail && asset.backend() == backend
            }));
        }
    }

    #[test]
    fn migraphx_required_files_include_accelerated_onnx_assets() {
        let files = required_model_files(ExecutionMode::MiGraphX);
        assert!(files.contains(&"segmentation-3.0-b32.onnx".to_string()));
        assert!(files.contains(&"wespeaker-fbank.onnx".to_string()));
        assert!(files.contains(&"wespeaker-fbank-b32.onnx".to_string()));
        assert!(files.contains(&"wespeaker-multimask-tail.onnx".to_string()));
        assert!(files.contains(&"wespeaker-multimask-tail-b32.onnx".to_string()));
        assert!(files.contains(&"wespeaker-voxceleb-resnet34-b64.onnx".to_string()));
    }
}