mozu-audio 0.3.0

//! Download and cache the compiled Core ML models used by Mozu Audio.

use std::{
    fs::{self, File, OpenOptions},
    io::{self, Read, Write},
    path::{Path, PathBuf},
    time::Duration,
};

use sha2::{Digest, Sha256};

use crate::{Cancellation, ComputeConfig, Error, HaasEngine, Result};

const HAAS_REVISION: &str = "e2f7a376a517494ba3903fb42bd70c89883c83ca";
const HAAS_15S_REVISION: &str = "e54549fdd81dadabd9776d51fb6f1b3a7180663f";
const HAAS_REPOSITORY: &str = "ooookao/haas-ja-v1.0-coreml";
const SILERO_REVISION: &str = "b419383c55c110e2c9271fa6ee0ea83d03c70d96";
const SILERO_MODEL: &str = "silero-vad-unified-256ms-v6.2.1.mlmodelc";

struct Bundle {
    repository: &'static str,
    revision: &'static str,
    directory: &'static str,
    file_sets: &'static [&'static [(&'static str, &'static str)]],
}

impl Bundle {
    fn files(&self) -> impl Iterator<Item = (&'static str, &'static str)> + '_ {
        self.file_sets.iter().copied().flatten().copied()
    }
}

const HAAS: Bundle = Bundle {
    repository: HAAS_REPOSITORY,
    revision: HAAS_REVISION,
    directory: "haas-ja-v1.0-int8",
    file_sets: &[HAAS_FILES, HAAS_PREDICTOR_FILES],
};
const HAAS_15S: Bundle = Bundle {
    repository: HAAS_REPOSITORY,
    revision: HAAS_15S_REVISION,
    directory: "haas-ja-v1.0-int8-1500",
    file_sets: &[HAAS_15S_FILES, HAAS_PREDICTOR_FILES],
};
const SILERO: Bundle = Bundle {
    repository: "FluidInference/silero-vad-coreml",
    revision: SILERO_REVISION,
    directory: "silero-vad-v6.2.1",
    file_sets: &[SILERO_FILES],
};

/// Maximum mel input length for a compiled Haas encoder.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum HaasVariant {
    Frames1500,
    Frames3072,
}

/// Bytes received for a file in the current download.
#[derive(Clone, Copy, Debug)]
pub struct DownloadProgress {
    pub file: &'static str,
    pub received_bytes: u64,
    pub total_bytes: Option<u64>,
}

/// Paths to the cached Haas assets, INT8 encoder and fused predictor.
#[derive(Clone, Debug)]
pub struct HaasModelFiles {
    pub assets: PathBuf,
    pub encoder: PathBuf,
    pub predictor: PathBuf,
}

impl HaasModelFiles {
    pub fn load(&self, compute: ComputeConfig) -> Result<HaasEngine> {
        HaasEngine::load_with_compute(&self.assets, &self.encoder, &self.predictor, compute)
    }
}

/// Stores versioned Core ML bundles in an application-selected directory.
pub struct ModelCache {
    root: PathBuf,
}

impl ModelCache {
    pub fn new(root: PathBuf) -> Self {
        Self { root }
    }

    /// Fetch Haas on the first call. A verified cache hit requires no network access.
    /// This operation blocks. Call it from a worker thread in a GUI application.
    pub fn prepare_haas(
        &self,
        cancellation: &Cancellation,
        progress: impl FnMut(DownloadProgress),
    ) -> Result<HaasModelFiles> {
        self.prepare_haas_variant(HaasVariant::Frames3072, cancellation, progress)
    }

    /// Fetch the selected Haas encoder and its matching assets.
    /// This operation blocks. Call it from a worker thread in GUI applications.
    pub fn prepare_haas_variant(
        &self,
        variant: HaasVariant,
        cancellation: &Cancellation,
        progress: impl FnMut(DownloadProgress),
    ) -> Result<HaasModelFiles> {
        let (bundle, reusable, assets, encoder) = match variant {
            HaasVariant::Frames1500 => (
                &HAAS_15S,
                &HAAS,
                "assets-1500",
                "encoder-int8-1500.mlmodelc",
            ),
            HaasVariant::Frames3072 => (&HAAS, &HAAS_15S, "assets", "encoder-int8.mlmodelc"),
        };
        let directory = self.prepare(bundle, Some(reusable), cancellation, progress)?;
        Ok(HaasModelFiles {
            assets: directory.join(assets),
            encoder: directory.join(encoder),
            predictor: directory.join("predictor.mlmodelc"),
        })
    }

    /// Fetch Silero VAD and return the path accepted by [crate::vad::SileroVad::load].
    /// This operation blocks. A verified cache hit requires no network access.
    pub fn prepare_silero(
        &self,
        cancellation: &Cancellation,
        progress: impl FnMut(DownloadProgress),
    ) -> Result<PathBuf> {
        self.prepare(&SILERO, None, cancellation, progress)
            .map(|directory| directory.join(SILERO_MODEL))
    }

    fn prepare(
        &self,
        bundle: &Bundle,
        reusable: Option<&Bundle>,
        cancellation: &Cancellation,
        mut progress: impl FnMut(DownloadProgress),
    ) -> Result<PathBuf> {
        cancellation.check()?;
        let parent = self.root.join(bundle.directory);
        io_at(&parent, fs::create_dir_all(&parent))?;
        let lock_path = parent.join("download.lock");
        let lock = io_at(
            &lock_path,
            OpenOptions::new()
                .create(true)
                .truncate(false)
                .read(true)
                .write(true)
                .open(&lock_path),
        )?;
        loop {
            cancellation.check()?;
            match lock.try_lock() {
                Ok(()) => break,
                Err(fs::TryLockError::WouldBlock) => std::thread::sleep(Duration::from_millis(50)),
                Err(fs::TryLockError::Error(source)) => {
                    return Err(Error::Io {
                        path: lock_path,
                        source,
                    });
                }
            }
        }
        let destination = parent.join(bundle.revision);
        if verified(&destination, bundle, cancellation)? {
            return Ok(destination);
        }

        let stage = io_at(&parent, tempfile::tempdir_in(&parent))?;
        let agent: ureq::Agent = ureq::Agent::config_builder()
            .timeout_global(Some(Duration::from_secs(600)))
            .build()
            .into();
        for (file, expected) in bundle.files() {
            cancellation.check()?;
            let path = stage.path().join(file);
            let folder = path
                .parent()
                .expect("model files are inside the staging directory");
            io_at(folder, fs::create_dir_all(folder))?;
            if let Some(reusable) = reusable
                && let Some((source, _)) = reusable.files().find(|&(_, digest)| digest == expected)
                && copy_cached_file(
                    &self
                        .root
                        .join(reusable.directory)
                        .join(reusable.revision)
                        .join(source),
                    &path,
                    expected,
                    cancellation,
                )?
            {
                continue;
            }
            let url = format!(
                "https://huggingface.co/{}/resolve/{}/{file}",
                bundle.repository, bundle.revision
            );
            let mut response = agent
                .get(&url)
                .call()
                .map_err(|source| Error::ModelDownload {
                    url: url.clone(),
                    message: source.to_string(),
                })?;
            let total_bytes = response.body().content_length();
            let mut writer = io_at(&path, File::create(&path))?;
            let digest = copy_and_hash(
                response.body_mut().as_reader(),
                &mut writer,
                &path,
                cancellation,
                |received_bytes| {
                    progress(DownloadProgress {
                        file,
                        received_bytes,
                        total_bytes,
                    })
                },
            )?;
            if digest != expected {
                return Err(Error::ModelDownload {
                    url,
                    message: format!("SHA-256 mismatch: expected={expected}, received={digest}"),
                });
            }
        }
        cancellation.check()?;
        match fs::remove_dir_all(&destination) {
            Ok(()) => {}
            Err(error) if error.kind() == io::ErrorKind::NotFound => {}
            Err(source) => {
                return Err(Error::Io {
                    path: destination,
                    source,
                });
            }
        }
        io_at(&destination, fs::rename(stage.path(), &destination))?;
        Ok(destination)
    }
}

fn verified(directory: &Path, bundle: &Bundle, cancellation: &Cancellation) -> Result<bool> {
    for (file, expected) in bundle.files() {
        let path = directory.join(file);
        let mut input = match File::open(&path) {
            Ok(input) => input,
            Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(false),
            Err(source) => return Err(Error::Io { path, source }),
        };
        let digest = copy_and_hash(&mut input, &mut io::sink(), &path, cancellation, |_| {})?;
        if digest != expected {
            return Ok(false);
        }
    }
    Ok(true)
}

fn copy_cached_file(
    source: &Path,
    destination: &Path,
    expected: &str,
    cancellation: &Cancellation,
) -> Result<bool> {
    let input = match File::open(source) {
        Ok(input) => input,
        Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(false),
        Err(error) => return io_at(source, Err(error)),
    };
    let output = io_at(destination, File::create(destination))?;
    let digest = copy_and_hash(input, output, destination, cancellation, |_| {})?;
    Ok(digest == expected)
}

fn copy_and_hash(
    mut input: impl Read,
    mut output: impl Write,
    path: &Path,
    cancellation: &Cancellation,
    mut progress: impl FnMut(u64),
) -> Result<String> {
    let mut hash = Sha256::new();
    let mut buffer = [0; 64 * 1024];
    let mut received = 0u64;
    loop {
        cancellation.check()?;
        let count = io_at(path, input.read(&mut buffer))?;
        if count == 0 {
            break;
        }
        let bytes = &buffer[..count];
        io_at(path, output.write_all(bytes))?;
        hash.update(bytes);
        received += count as u64;
        progress(received);
    }
    Ok(hash
        .finalize()
        .iter()
        .map(|byte| format!("{byte:02x}"))
        .collect())
}

fn io_at<T>(path: &Path, result: io::Result<T>) -> Result<T> {
    result.map_err(|source| Error::Io {
        path: path.to_path_buf(),
        source,
    })
}

const HAAS_FILES: &[(&str, &str)] = &[
    (
        "LICENSE",
        "cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30",
    ),
    (
        "NOTICE",
        "7caad652ba44030bfb8c72962433874d481924328692e5f3cf5a58c0a9f8ccca",
    ),
    (
        "assets/embeddings.f32",
        "3eb3df5cfc5efc5dee3eedb1dfe0be7daf03bca4e660d9ccc86aa6775050a433",
    ),
    (
        "assets/filterbank.f32",
        "eb5cf9eb93b8038062ebc849b14ec85098abe1c915bc22262be8ad336bcf563a",
    ),
    (
        "assets/metadata.json",
        "09211050a5142b13274b85b47de3beb567e6bf478a52dbc5273273524aa413f6",
    ),
    (
        "assets/tokenizer.json",
        "aa665c2c957db3e4ab8daba78b03779bcf87b75b8f084fdcebf13e8793bd10d0",
    ),
    (
        "assets/window.f32",
        "6a66c2c694505b41b8987c4885a618e1a32612b43a4935a208a3c9c956b22d02",
    ),
    (
        "encoder-int8.mlmodelc/analytics/coremldata.bin",
        "0c4787ed4f560c2bf4ab4bfdac691c4b65580b9bb52980d9d5dbf328f63b955b",
    ),
    (
        "encoder-int8.mlmodelc/coremldata.bin",
        "7f0faeb1ace3d85a7a2886c5661e32e9b2f7a5ebd3d4bf601b5e35a021a7fcdb",
    ),
    (
        "encoder-int8.mlmodelc/model.mil",
        "b8cd70086832633bd987f5cb4c9bf5182890d1c05b88493e6a08b9d318ec56f8",
    ),
    (
        "encoder-int8.mlmodelc/weights/weight.bin",
        "6e3df99334be3af5a704d28d3dadcdf08b60eb51b0350f162c7063f7312462df",
    ),
];
const HAAS_15S_FILES: &[(&str, &str)] = &[
    (
        "assets-1500/embeddings.f32",
        "3eb3df5cfc5efc5dee3eedb1dfe0be7daf03bca4e660d9ccc86aa6775050a433",
    ),
    (
        "assets-1500/filterbank.f32",
        "eb5cf9eb93b8038062ebc849b14ec85098abe1c915bc22262be8ad336bcf563a",
    ),
    (
        "assets-1500/metadata.json",
        "7afbd56913aa6bba2d2aa80f7f323b63dccfbddaa2908782f531acc857a4fb3b",
    ),
    (
        "assets-1500/tokenizer.json",
        "aa665c2c957db3e4ab8daba78b03779bcf87b75b8f084fdcebf13e8793bd10d0",
    ),
    (
        "assets-1500/window.f32",
        "6a66c2c694505b41b8987c4885a618e1a32612b43a4935a208a3c9c956b22d02",
    ),
    (
        "encoder-int8-1500.mlmodelc/analytics/coremldata.bin",
        "d9fd93ba4c586634572a753b6877b6d7bb948be73638a3bb48b9675e7c51de35",
    ),
    (
        "encoder-int8-1500.mlmodelc/coremldata.bin",
        "499563166f4c862389fafc7bae0bd717b255890edd3046b8caa8aa91bdf1b18f",
    ),
    (
        "encoder-int8-1500.mlmodelc/model.mil",
        "88a7089c5d496dcbfca87a050ce8e20594c532e9a2578e9cb956c0ee11370fe2",
    ),
    (
        "encoder-int8-1500.mlmodelc/weights/weight.bin",
        "dc4d4529128429e9a8f997a027de98c1bd4add721ba5052b7df9f7801431cab0",
    ),
];
const HAAS_PREDICTOR_FILES: &[(&str, &str)] = &[
    (
        "predictor.mlmodelc/analytics/coremldata.bin",
        "2734406335ca07481a8ddb822ff051aad49099d1ae2b421d59f090c204cb00c0",
    ),
    (
        "predictor.mlmodelc/coremldata.bin",
        "d9ca21a95faec3241a70d437ae311fd63666905f12465b4d722728402d73e42f",
    ),
    (
        "predictor.mlmodelc/model.mil",
        "16be0867bd53d1cb15559889e9b0ad6f9122be1275fd8a2c4b09ce1f99b4fc20",
    ),
    (
        "predictor.mlmodelc/weights/weight.bin",
        "59620240b377cf1f777d9c7a0acc8406692035e8341656279a9633cbecc134c4",
    ),
    (
        "tokenizer.model",
        "9bd7bcdff62e43200001cd4c258d22e475b7b3b1e0f64f8be8e6e046f402f38c",
    ),
];

const SILERO_FILES: &[(&str, &str)] = &[
    (
        "silero-vad-unified-256ms-v6.2.1.mlmodelc/analytics/coremldata.bin",
        "8067594eb3126ab8318af507f0c00cabfed40d5fedb8a0ee5075dd02e903d909",
    ),
    (
        "silero-vad-unified-256ms-v6.2.1.mlmodelc/coremldata.bin",
        "7db35a4fd995222a7fb0129713473b15d1462572ab4a2e5e4d56bcaad9e40f41",
    ),
    (
        "silero-vad-unified-256ms-v6.2.1.mlmodelc/metadata.json",
        "2740be542c611e1ba358e1849b4e265c65cdf0b17192767e1e5de86a31ac94d6",
    ),
    (
        "silero-vad-unified-256ms-v6.2.1.mlmodelc/model.mil",
        "c6a9d1bf22d413265da0a07a1d14151c3ea2fad296b3aa5859275b33ef1c3270",
    ),
    (
        "silero-vad-unified-256ms-v6.2.1.mlmodelc/weights/weight.bin",
        "53ecc8b5081146140ab654c89109cf001f2183abddd7a2411c5081feeffff063",
    ),
];