diff --git a/Documentation/Models.md b/Documentation/Models.md index 3aea8b7f8..c0061337d 100644 --- a/Documentation/Models.md +++ b/Documentation/Models.md @@ -58,6 +58,7 @@ TDT/CTC and the non-autoregressive models above are wrapped by `SlidingWindowAsr | **Kokoro ANE (7-stage)** | Kokoro 82M weights split into 7 CoreML stages so the ANE-friendly layers (Albert / Prosody / Vocoder) stay resident on the Neural Engine while PostAlbert / Alignment / Noise / Tail run on CPU. 3-11× RTFx. English (`af_heart`) and Mandarin (`ANE-zh`) variants. ≤510 IPA phonemes per call, no chunker / SSML / custom lexicon. Managed by `KokoroAneManager`. | ANE-optimized variant derived (with permission) from [laishere/kokoro-coreml](https://github.com/laishere/kokoro-coreml). The original single-graph (mono) Kokoro backend was removed in favor of this ANE pipeline; the kokoro repo root is now retained only for shared G2P assets. | | **PocketTTS** | TTS backend (~155M params). Autoregressive frame-by-frame generation with dynamic audio chunking. No phoneme stage, works directly on text tokens. Managed by `PocketTtsManager`. | Supports streaming, minimal RAM usage, excellent quality | | **Supertonic-3** | Multilingual TTS, 31 languages (`en`, `ko`, `ja`, `ar`, `bg`, `cs`, `da`, `de`, `el`, `es`, `et`, `fi`, `fr`, `hi`, `hr`, `hu`, `id`, `it`, `lt`, `lv`, `nl`, `pl`, `pt`, `ro`, `ru`, `sk`, `sl`, `sv`, `tr`, `uk`, `vi`, plus `na` for numeric/language-agnostic input). 4-stage CoreML pipeline (text_encoder → duration_predictor → vector_estimator → vocoder, ~398 MB). Caller-supplied voice styles loaded from Supertonic preset JSON. 44.1 kHz mono fp32 output. Managed by `Supertonic3Manager`. | CoreML conversion of upstream `Supertone/supertonic-3`; see `Scripts/convert_supertonic3_to_coreml.py`. | +| **MOSS-TTS-Nano (beta)** | 0.1B multilingual streaming TTS with zero-shot voice cloning (20 languages incl. `zh`, `en`, `de`, `es`, `fr`, `ja`, `ko`, `ru`, `ar`). GPT-2 global LM + local transformer over MOSS-Audio-Tokenizer-Nano codes, decoded per 80 ms frame by a streaming codec (`Prefill → Frame → CodecStep → Step`). Voices are codec-token sequences of a reference clip (two presets, or `cloneVoice(audioURL:)`). Native 48 kHz **stereo** fp32 output. Managed by `MossTtsNanoManager`. | CoreML conversion of upstream `OpenMOSS-Team/MOSS-TTS-Nano-100M`; see mobius `models/tts/moss-tts-nano/coreml` and `Documentation/TTS/MossTtsNano.md`. | | **StyleTTS2 (LibriTTS, iteration_3)** | Reference-audio–driven zero-shot English TTS. 8-stage CoreML pipeline (`text_encoder → bert → ref_encoder → fused_diffusion_sampler → duration_predictor → fused_f0n_har_source → decoder_pre → decoder_upsample`) with 3 lazily-loaded T = 64 / 128 / 256 bucket variants of `bert` / `fused_diffusion_sampler`. 5-step ADPM2 Karras-σ diffusion sampler with α/β style blending against a speaker reference clip. 24 kHz mono fp32 output. Phonemizer reuses Kokoro's Misaki lexicon cache + BART G2P CoreML model with Misaki uppercase diphthong shorthand (`A O I Y W` → `eɪ oʊ aɪ ɔɪ aʊ`) expanded before encoding so the output matches the espeak IPA the model was trained on. Callers with a higher-quality phonemizer can bypass the stack via `StyleTTS2Manager.synthesize(ipa:...)`. See [StyleTTS2.md](TTS/StyleTTS2.md). | Zero-shot voice cloning from a single reference WAV; English only | ## Evaluated Models (Not Supported) @@ -93,6 +94,7 @@ Models we converted and tested but are not supported: too large for on-device de | PocketTTS | [FluidInference/pocket-tts-coreml](https://huggingface.co/FluidInference/pocket-tts-coreml) | | StyleTTS2 (LibriTTS, iteration_3) | [FluidInference/StyleTTS-2-coreml/iteration_3/compiled](https://huggingface.co/FluidInference/StyleTTS-2-coreml/tree/main/iteration_3/compiled) (shared phonemizer assets pulled from [`FluidInference/kokoro-82m-coreml`](https://huggingface.co/FluidInference/kokoro-82m-coreml): `G2PEncoder.mlmodelc`, `G2PDecoder.mlmodelc`, `g2p_vocab.json`, `us_lexicon_cache.json`) | | Supertonic-3 | [FluidInference/supertonic-3-coreml](https://huggingface.co/FluidInference/supertonic-3-coreml) | +| MOSS-TTS-Nano (beta) | [FluidInference/moss-tts-nano-coreml](https://huggingface.co/FluidInference/moss-tts-nano-coreml) (`voices/*.json` presets, fp32 `MossNano-CodecEncoder-fp32.mlmodelc` fetched only for custom cloning) | | Nemotron Streaming (English) | [FluidInference/nemotron-speech-streaming-en-0.6b-coreml](https://huggingface.co/FluidInference/nemotron-speech-streaming-en-0.6b-coreml) (subdirs: `/560ms`, `/1120ms`, `/2240ms`) | | Nemotron Streaming (Multilingual) | [FluidInference/Nemotron-3.5-ASR-Streaming-Multilingual-0.6b-CoreML](https://huggingface.co/FluidInference/Nemotron-3.5-ASR-Streaming-Multilingual-0.6b-CoreML) (`/ms/` subfolders) | | Multilingual G2P (Charsiu ByT5) | [FluidInference/charsiu-g2p-byt5-coreml](https://huggingface.co/FluidInference/charsiu-g2p-byt5-coreml) | diff --git a/Documentation/TTS/MossTtsNano.md b/Documentation/TTS/MossTtsNano.md new file mode 100644 index 000000000..3e92004a5 --- /dev/null +++ b/Documentation/TTS/MossTtsNano.md @@ -0,0 +1,141 @@ +# MOSS-TTS-Nano Swift Inference + +> **Beta.** This is a beta model conversion: the Swift API, the published model +> artifacts and the output quality may change between releases. + +0.1B multilingual streaming TTS with zero-shot voice cloning. 20 languages, native +48 kHz **stereo** Float32 output, 80 ms frames streamed as they are decoded. +Six CoreML bundles (≈ 275 MB fp16 on disk; the fp32 voice encoder is fetched only +for custom cloning). + +## Overview + +MOSS-TTS-Nano (OpenMOSS, Apache-2.0) is a pure autoregressive "audio tokenizer + +LLM" model: a 12-layer GPT-2 (768-d, RoPE) reads rows of `[text token, 16 codec +codes]`, a 1-layer local transformer emits the 16 RVQ codebooks of each 12.5 Hz +frame, and MOSS-Audio-Tokenizer-Nano (22M, causal transformer codec) turns frames +into 48 kHz stereo audio. The conversion lives in +[mobius `models/tts/moss-tts-nano/coreml`](https://github.com/FluidInference/mobius/tree/main/models/tts/moss-tts-nano/coreml); +weights are published at +[FluidInference/moss-tts-nano-coreml](https://huggingface.co/FluidInference/moss-tts-nano-coreml). + +``` +reference clip ─CodecEncoder (fp32)─► 16×T codes ─┐ +text ─SentencePiece BPE─► ids ────────────────────┴─► rows [T,17] ─Prefill─► hidden + KV + │ per 80 ms frame + Frame (local transformer + sampler) ─► 16 codes ─► CodecStep ─► 3840 stereo samples + Step (global GPT-2, KV update) ◄── [assistant_slot, 16 codes] +``` + +Sampling (top-k / top-p / temperature / repetition penalty) is inside the `Frame` +graph; the host only supplies uniform randoms, so a `seed` makes a run +reproducible on the same machine. + +## Quick Start + +### CLI + +```bash +# Preset voice (en_2), stereo 48 kHz WAV +swift run fluidaudiocli tts "Hello from the neural engine." \ + --backend moss-tts-nano --output hello.wav + +# Clone a voice from a clip (≤ ~25 s), save the codes, reuse them later +swift run fluidaudiocli tts "Same voice, new words." --backend moss \ + --clone-voice speaker.wav --save-voice speaker.json --output clone.wav +swift run fluidaudiocli tts "Again." --backend moss --voice-file speaker.json + +# Deterministic sampling; tokenizer parity check +swift run fluidaudiocli tts "Hello" --backend moss --seed 7 --metrics moss.json +swift run fluidaudiocli tts "Hello" --backend moss --tokens-only +``` + +| Flag | Default | Notes | +|---|---|---| +| `--voice ` | `en_2` | Preset: `en_2` (English), `zh_1` (Mandarin) | +| `--clone-voice ` | – | Any AVFoundation-readable file; resampled to 48 kHz stereo | +| `--save-voice` / `--voice-file ` | – | Persist / reuse cloned voice codes | +| `--seed N` | system RNG | Seed for the in-graph sampler | +| `--greedy` | off | Argmax decoding (parity oracle; upstream greedy never stops) | +| `--cpu-only` | off | Everything on CPU | + +### Swift + +```swift +import FluidAudio + +let manager = try await MossTtsNanoManager.downloadAndCreate() +let voice = try await manager.loadVoice(.en2) // or manager.cloneVoice(audioURL:) + +// Batch +let audio = try await manager.synthesize(text: "Local speech synthesis.", voice: voice) +play(left: audio.left, right: audio.right, sampleRate: audio.sampleRate) // 48 kHz +// audio.mono / audio.interleaved for single-channel or interleaved sinks + +// Streaming — 80 ms stereo frames as they are decoded +var options = MossTtsNanoSamplingOptions() +options.seed = 42 +for try await frame in try await manager.synthesizeStreaming(text: text, voice: voice, options: options) { + schedule(left: frame.left, right: frame.right) // frame.isPause marks inter-chunk silence +} +``` + +`MossTtsNanoSamplingOptions` exposes the upstream defaults (`textTemperature` 1.5, +`audioTemperature` 1.7, `audioTopP` 0.8, `repetitionPenalty` 1.0, `maxNewFrames` +375 = 30 s per chunk, `maxTextTokens` 50). + +## Text handling + +- Tokenizer: SentencePiece **BPE** (`MossTtsNanoTokenizer`) with `nmt_nfkc` + normalization, dummy prefix and byte fallback — verified id-for-id against the + upstream `sentencepiece` processor. Upstream additionally runs WeTextProcessing + (number/date verbalization); that step is not ported, so spell out numbers or + run `NemoTextNormalizer` first for best results. +- Chunking: upstream's sentence → clause → token-budget splitter + (`MossTtsNanoTextChunker`), 50 tokens per chunk, first letter capitalized and a + terminal period added when missing; chunks are joined with 0.40 s (≤ 4 words) or + 0.24 s pauses. +- Prompt: `user … Reference: [reference rows] … + Text: {ids} assistant `; the template + ids and special tokens come from the repo's `config.json`. + +## Voices + +A voice is the codec token sequence of a reference clip (`MossTtsNanoVoice`, +`[frames][16]`). Two presets are published; `cloneVoice(audioURL:)` encodes any +clip with the fp32 codec encoder (fp16 loses a third of the codes through the +residual quantizer). The prefill graph holds 512 rows, so reference frames + text +tokens + 78 template rows must fit: keep clips under ~25 s. + +## Performance (M5 Pro, macOS 26.7, warm) + +| Stage | ms / call | Unit | +|---|---|---| +| Prefill (512 rows) | 11 | GPU | +| Step (KV M=1024) | 7.6 | GPU | +| Frame | 5.3 | any | +| CodecStep | 5.2 | GPU / ANE | +| CodecEncoder fp32 (8 s clip) | 22 | GPU | + +≈ 18 ms of model compute per 80 ms frame. End to end through the Swift host +(release CLI, `en_2` voice, 9–10 s utterances): **3.1–3.8× real time**, first audio +0.22–0.40 s after the call (includes prefill and the first frame); a 3 s Mandarin +utterance with `zh_1` runs at 2.5×. Prefill and Step fail ANE compilation (`ANECCompile FAILED`) and are +pinned to CPU+GPU by `MossTtsNanoModelStore`; `computeUnits` applies to Frame and +CodecStep. The Step graph round-trips a 38 MB KV cache per frame; a `StateType` +(iOS 18) variant is a planned follow-up. + +## Parity + +- Wrappers vs upstream fp32: hidden states 7e-6, codec 234 dB SNR, encoder codes exact. +- CoreML fp16 greedy replay of a 375-frame upstream reference: 370/375 frames token-exact. +- Streaming codec step vs full decode: 56.6 dB SNR (GPU). +- Parakeet ASR, two English phrases with the `en_2` voice: 8.3 % WER (CoreML) vs + 10.1 % (upstream PyTorch fp32). + +## Known limitations + +- No text normalization (numbers, dates) — see above. +- `nq < 16` low-bitrate decoding and the batch `CodecDecoder` are not exposed in Swift. +- Sentence chunks are independent generations (upstream behaviour); prosody is not + carried across chunk boundaries. diff --git a/Sources/FluidAudio/ModelNames.swift b/Sources/FluidAudio/ModelNames.swift index a5ffc4c43..e0aeed8b2 100644 --- a/Sources/FluidAudio/ModelNames.swift +++ b/Sources/FluidAudio/ModelNames.swift @@ -77,6 +77,12 @@ public enum Repo: String, CaseIterable, Sendable { /// recipe. Ships four `.mlmodelc` bundles + `tts.json` + /// `unicode_indexer.json` at the repo root. case supertonic3 = "FluidInference/supertonic-3-coreml" + /// MOSS-TTS-Nano (OpenMOSS) — 0.1B multilingual streaming TTS with zero-shot + /// voice cloning, 48 kHz stereo. Four streaming-path `.mlmodelc` bundles + /// (Prefill / Step / Frame / CodecStep) + `config.json` + `tokenizer.model` at + /// the repo root, fp32 `CodecEncoder` for voice cloning, `voices/*.json` presets. + /// Conversion: mobius `models/tts/moss-tts-nano/coreml`. + case mossTtsNano = "FluidInference/moss-tts-nano-coreml" /// NeuTTS-2E emotional English TTS (Qwen3 236M backbone + NeuCodec /// decoder). Compiled `.mlmodelc` bundles + `tokenizer.json` + /// `samples/.json` reference codes at the repo root; the @@ -178,6 +184,8 @@ public enum Repo: String, CaseIterable, Sendable { return "inflect-v2-coreml/micro" case .inflectNano: return "inflect-v2-coreml/nano" + case .mossTtsNano: + return "moss-tts-nano-coreml" } } @@ -1204,6 +1212,34 @@ public enum ModelNames { /// Supertonic-3 multilingual TTS — 4 `.mlmodelc` bundles + 2 companion /// JSON files. File names match the HuggingFace tree at /// `FluidInference/supertonic-3-coreml/`. + public enum MossTtsNano { + public static let prefill = "MossNano-Prefill-T512-M1024-fp16" + public static let step = "MossNano-Step-M1024-fp16" + public static let frame = "MossNano-Frame-fp16" + public static let codecStep = "MossNano-CodecStep-fp16" + public static let codecDecoder = "MossNano-CodecDecoder-fp16" + public static let codecEncoder = "MossNano-CodecEncoder-fp32" + + public static let prefillFile = prefill + ".mlmodelc" + public static let stepFile = step + ".mlmodelc" + public static let frameFile = frame + ".mlmodelc" + public static let codecStepFile = codecStep + ".mlmodelc" + /// Flexible-length batch decoder (≤ 125 frames); not used by the Swift streaming path. + public static let codecDecoderFile = codecDecoder + ".mlmodelc" + /// fp32 prompt encoder for custom voice cloning; downloaded on demand. + public static let codecEncoderFile = codecEncoder + ".mlmodelc" + + public static let configFile = "config.json" + public static let tokenizerFile = "tokenizer.model" + public static let voicesSubdir = "voices" + + /// Streaming-path bundles loaded by `MossTtsNanoModelStore`. + public static let requiredModels: Set = [prefillFile, stepFile, frameFile, codecStepFile] + + /// Models + companion files fetched on first use. + public static let requiredFiles: Set = requiredModels.union([configFile, tokenizerFile]) + } + public enum Supertonic3 { public static let textEncoder = "TextEncoder" public static let durationPredictor = "DurationPredictor" @@ -1630,6 +1666,8 @@ public enum ModelNames { return ModelNames.LuxTts.requiredFiles(variant: variant) case .inflectMicro, .inflectNano: return ModelNames.Inflect.requiredModels + case .mossTtsNano: + return ModelNames.MossTtsNano.requiredFiles } } } diff --git a/Sources/FluidAudio/Shared/Download/ModelHub.swift b/Sources/FluidAudio/Shared/Download/ModelHub.swift index 0df9c9fcd..521e318de 100644 --- a/Sources/FluidAudio/Shared/Download/ModelHub.swift +++ b/Sources/FluidAudio/Shared/Download/ModelHub.swift @@ -770,7 +770,9 @@ public enum ModelHub { itemPath.hasSuffix(".json") || itemPath.hasSuffix(".model") || itemPath.hasSuffix(".bin") return isInSubPath && (matchesPattern || isMetadata) } - return patterns.isEmpty || patterns.contains { itemPath.hasPrefix($0) } + // Patterns carry a trailing "/" (bundle directories); a required root-level + // *file* (e.g. `tokenizer.model`) matches when the pattern is exactly its path. + return patterns.isEmpty || patterns.contains { itemPath.hasPrefix($0) || $0 == itemPath + "/" } || itemPath.hasSuffix(".json") || itemPath.hasSuffix(".txt") } } diff --git a/Sources/FluidAudio/TTS/MossTtsNano/Assets/MossTtsNanoAudioLoader.swift b/Sources/FluidAudio/TTS/MossTtsNano/Assets/MossTtsNanoAudioLoader.swift new file mode 100644 index 000000000..3d35d2d96 --- /dev/null +++ b/Sources/FluidAudio/TTS/MossTtsNano/Assets/MossTtsNanoAudioLoader.swift @@ -0,0 +1,63 @@ +@preconcurrency import AVFoundation +import Foundation +import os + +/// Decodes any AVFoundation-readable file to 48 kHz stereo Float32 (mono is duplicated). +enum MossTtsNanoAudioLoader { + + static func loadStereo48k(url: URL) throws -> (left: [Float], right: [Float]) { + do { + let file = try AVAudioFile(forReading: url) + guard + let target = AVAudioFormat( + commonFormat: .pcmFormatFloat32, sampleRate: Double(MossTtsNanoConstants.sampleRate), + channels: 2, interleaved: false) + else { throw MossTtsNanoError.audioLoadFailed(path: url.path, underlying: "cannot build 48 kHz format") } + let frameCount = AVAudioFrameCount(file.length) + guard let input = AVAudioPCMBuffer(pcmFormat: file.processingFormat, frameCapacity: max(frameCount, 1)) + else { throw MossTtsNanoError.audioLoadFailed(path: url.path, underlying: "cannot allocate buffer") } + try file.read(into: input) + + guard let converter = AVAudioConverter(from: file.processingFormat, to: target) else { + throw MossTtsNanoError.audioLoadFailed( + path: url.path, underlying: "no converter for \(file.processingFormat)") + } + let ratio = target.sampleRate / file.processingFormat.sampleRate + let capacity = AVAudioFrameCount(Double(input.frameLength) * ratio) + 4096 + guard let output = AVAudioPCMBuffer(pcmFormat: target, frameCapacity: capacity) else { + throw MossTtsNanoError.audioLoadFailed(path: url.path, underlying: "cannot allocate output buffer") + } + let provided = OSAllocatedUnfairLock(initialState: false) + let inputBlock: AVAudioConverterInputBlock = { _, outStatus in + let wasProvided = provided.withLock { state -> Bool in + if state { return true } + state = true + return false + } + if wasProvided { + outStatus.pointee = .endOfStream + return nil + } + outStatus.pointee = .haveData + return input + } + var conversionError: NSError? + let status = converter.convert(to: output, error: &conversionError, withInputFrom: inputBlock) + if status == .error || conversionError != nil { + throw MossTtsNanoError.audioLoadFailed( + path: url.path, underlying: conversionError?.localizedDescription ?? "conversion failed") + } + let n = Int(output.frameLength) + guard let channels = output.floatChannelData, n > 0 else { + throw MossTtsNanoError.audioLoadFailed(path: url.path, underlying: "empty audio") + } + let left = Array(UnsafeBufferPointer(start: channels[0], count: n)) + let right = Array(UnsafeBufferPointer(start: channels[1], count: n)) + return (left, right) + } catch let error as MossTtsNanoError { + throw error + } catch { + throw MossTtsNanoError.audioLoadFailed(path: url.path, underlying: "\(error)") + } + } +} diff --git a/Sources/FluidAudio/TTS/MossTtsNano/Assets/MossTtsNanoModelStore.swift b/Sources/FluidAudio/TTS/MossTtsNano/Assets/MossTtsNanoModelStore.swift new file mode 100644 index 000000000..5ed1d057c --- /dev/null +++ b/Sources/FluidAudio/TTS/MossTtsNano/Assets/MossTtsNanoModelStore.swift @@ -0,0 +1,111 @@ +@preconcurrency import CoreML +import Foundation + +/// Holds the four streaming-path CoreML models plus config and tokenizer. +/// +/// Prefill and Step are pinned to CPU+GPU: their graphs fail ANE compilation +/// (`ANECCompile FAILED`) and would only fall back after a slow compile attempt. +/// Frame and CodecStep run on any unit; the caller's `computeUnits` applies to +/// them. The fp32 codec encoder (voice cloning) is loaded on demand. +public actor MossTtsNanoModelStore { + + private let logger = AppLogger(category: "MossTtsNanoModelStore") + + private let directory: URL? + private let computeUnits: MLComputeUnits + + private var repoDirectory: URL? + private var prefillModel: MLModel? + private var stepModel: MLModel? + private var frameModel: MLModel? + private var codecStepModel: MLModel? + private var encoderModel: MLModel? + private var loadedConfig: MossTtsNanoConfig? + private var loadedTokenizer: MossTtsNanoTokenizer? + + public init(directory: URL? = nil, computeUnits: MLComputeUnits = .cpuAndGPU) { + self.directory = directory + self.computeUnits = computeUnits + } + + private var lmComputeUnits: MLComputeUnits { + computeUnits == .cpuOnly ? .cpuOnly : .cpuAndGPU + } + + public func loadIfNeeded() async throws { + if prefillModel != nil { return } + let repoDir = try await MossTtsNanoResourceDownloader.ensureModels(directory: directory) + repoDirectory = repoDir + + do { + loadedConfig = try MossTtsNanoConfig.load( + from: repoDir.appendingPathComponent(ModelNames.MossTtsNano.configFile)) + } catch { + throw MossTtsNanoError.configLoadFailed("\(error)") + } + loadedTokenizer = try MossTtsNanoTokenizer( + modelURL: repoDir.appendingPathComponent(ModelNames.MossTtsNano.tokenizerFile)) + + logger.info("Loading MOSS-TTS-Nano CoreML models from \(repoDir.path)…") + let start = Date() + let lmConfig = MLModelConfiguration() + lmConfig.computeUnits = lmComputeUnits + let anyConfig = MLModelConfiguration() + anyConfig.computeUnits = computeUnits + + prefillModel = try load(repoDir: repoDir, fileName: ModelNames.MossTtsNano.prefillFile, config: lmConfig) + stepModel = try load(repoDir: repoDir, fileName: ModelNames.MossTtsNano.stepFile, config: lmConfig) + frameModel = try load(repoDir: repoDir, fileName: ModelNames.MossTtsNano.frameFile, config: anyConfig) + codecStepModel = try load(repoDir: repoDir, fileName: ModelNames.MossTtsNano.codecStepFile, config: anyConfig) + logger.info("MOSS-TTS-Nano models loaded in \(String(format: "%.2f", Date().timeIntervalSince(start)))s") + } + + // MARK: - Accessors + + public func prefill() throws -> MLModel { try unwrap(prefillModel) } + public func step() throws -> MLModel { try unwrap(stepModel) } + public func frame() throws -> MLModel { try unwrap(frameModel) } + public func codecStep() throws -> MLModel { try unwrap(codecStepModel) } + public func config() throws -> MossTtsNanoConfig { try unwrap(loadedConfig) } + public func tokenizer() throws -> MossTtsNanoTokenizer { try unwrap(loadedTokenizer) } + + /// fp32 codec encoder, downloaded and loaded on first call. + public func encoder() async throws -> MLModel { + if let encoderModel { return encoderModel } + let url = try await MossTtsNanoResourceDownloader.ensureEncoder(directory: directory) + let cfg = MLModelConfiguration() + cfg.computeUnits = lmComputeUnits + let model = try load(repoDir: url.deletingLastPathComponent(), fileName: url.lastPathComponent, config: cfg) + encoderModel = model + return model + } + + public func unload() { + prefillModel = nil + stepModel = nil + frameModel = nil + codecStepModel = nil + encoderModel = nil + } + + // MARK: - Helpers + + private func unwrap(_ value: T?) throws -> T { + guard let value else { throw MossTtsNanoError.notInitialized } + return value + } + + private func load(repoDir: URL, fileName: String, config: MLModelConfiguration) throws -> MLModel { + let url = repoDir.appendingPathComponent(fileName) + guard FileManager.default.fileExists(atPath: url.path) else { + throw MossTtsNanoError.modelFileNotFound(fileName) + } + do { + let model = try MLModel(contentsOf: url, configuration: config) + logger.info("Loaded \(fileName)") + return model + } catch { + throw MossTtsNanoError.corruptedModel(fileName, underlying: "\(error)") + } + } +} diff --git a/Sources/FluidAudio/TTS/MossTtsNano/Assets/MossTtsNanoResourceDownloader.swift b/Sources/FluidAudio/TTS/MossTtsNano/Assets/MossTtsNanoResourceDownloader.swift new file mode 100644 index 000000000..7862abb1f --- /dev/null +++ b/Sources/FluidAudio/TTS/MossTtsNano/Assets/MossTtsNanoResourceDownloader.swift @@ -0,0 +1,106 @@ +import Foundation + +/// Downloads the MOSS-TTS-Nano CoreML assets from `FluidInference/moss-tts-nano-coreml`. +/// +/// Required on first use: `Prefill`, `Step`, `Frame`, `CodecStep` bundles plus +/// `config.json` and `tokenizer.model`. The fp32 `CodecEncoder` is fetched lazily +/// by `MossTtsNanoModelStore.encoder()` (custom voice cloning only), and built-in +/// voices come from `voices/.json`. +public enum MossTtsNanoResourceDownloader { + + private static let logger = AppLogger(category: "MossTtsNanoResourceDownloader") + + @discardableResult + public static func ensureModels( + directory: URL? = nil, + progressHandler: ProgressHandler? = nil + ) async throws -> URL { + let modelsRoot = try directory ?? defaultCacheRoot() + let repoDir = modelsRoot.appendingPathComponent(Repo.mossTtsNano.folderName) + let missing = ModelNames.MossTtsNano.requiredFiles.filter { + !FileManager.default.fileExists(atPath: repoDir.appendingPathComponent($0).path) + } + if !missing.isEmpty { + logger.info("Downloading MOSS-TTS-Nano CoreML assets from HuggingFace (\(missing.count) missing)…") + do { + try await ModelHub.download(.mossTtsNano, to: modelsRoot, progressHandler: progressHandler) + } catch { + throw MossTtsNanoError.downloadFailed("\(error)") + } + } else { + logger.info("MOSS-TTS-Nano assets found in cache at \(repoDir.path)") + } + return repoDir + } + + /// Fetch the fp32 codec encoder used to turn a reference clip into voice codes. + @discardableResult + public static func ensureEncoder( + directory: URL? = nil, + progressHandler: ProgressHandler? = nil + ) async throws -> URL { + let modelsRoot = try directory ?? defaultCacheRoot() + let repoDir = modelsRoot.appendingPathComponent(Repo.mossTtsNano.folderName) + let encoderURL = repoDir.appendingPathComponent(ModelNames.MossTtsNano.codecEncoderFile) + if FileManager.default.fileExists(atPath: encoderURL.path) { return encoderURL } + logger.info("Downloading MOSS-TTS-Nano codec encoder from HuggingFace…") + do { + try await ModelHub.download( + .mossTtsNano, to: modelsRoot, + additionalModelNames: [ModelNames.MossTtsNano.codecEncoderFile], + progressHandler: progressHandler) + } catch { + throw MossTtsNanoError.downloadFailed("codec encoder: \(error)") + } + guard FileManager.default.fileExists(atPath: encoderURL.path) else { + throw MossTtsNanoError.downloadFailed("codec encoder missing after download") + } + return encoderURL + } + + /// Download (if needed) a built-in voice JSON and return its local URL. + @discardableResult + public static func ensureVoice( + _ voice: MossTtsNanoBuiltInVoice, + directory: URL? = nil, + progressHandler: ProgressHandler? = nil + ) async throws -> URL { + let modelsRoot = try directory ?? defaultCacheRoot() + let repoDir = modelsRoot.appendingPathComponent(Repo.mossTtsNano.folderName) + let localURL = repoDir.appendingPathComponent(voice.fileName) + if FileManager.default.fileExists(atPath: localURL.path) { return localURL } + logger.info("Downloading MOSS-TTS-Nano voice \(voice.rawValue) from HuggingFace…") + do { + try await ModelHub.download( + .mossTtsNano, + subdirectory: MossTtsNanoBuiltInVoiceFiles.subdirectory, + to: repoDir, + progressHandler: progressHandler, + shouldSkip: { $0 != voice.fileName } + ) + } catch { + throw MossTtsNanoError.downloadFailed("voice \(voice.rawValue): \(error)") + } + guard FileManager.default.fileExists(atPath: localURL.path) else { + throw MossTtsNanoError.downloadFailed("voice \(voice.rawValue) missing after download") + } + return localURL + } + + public static func loadVoice( + _ voice: MossTtsNanoBuiltInVoice, + directory: URL? = nil, + progressHandler: ProgressHandler? = nil + ) async throws -> MossTtsNanoVoice { + let url = try await ensureVoice(voice, directory: directory, progressHandler: progressHandler) + return try MossTtsNanoVoice.load(from: url) + } + + static func defaultCacheRoot() throws -> URL { + let root = try TtsCacheDirectory.ensure().appendingPathComponent("Models") + if !FileManager.default.fileExists(atPath: root.path) { + try FileManager.default.createDirectory(at: root, withIntermediateDirectories: true) + } + return root + } +} diff --git a/Sources/FluidAudio/TTS/MossTtsNano/MossTtsNanoConstants.swift b/Sources/FluidAudio/TTS/MossTtsNano/MossTtsNanoConstants.swift new file mode 100644 index 000000000..c484b42df --- /dev/null +++ b/Sources/FluidAudio/TTS/MossTtsNano/MossTtsNanoConstants.swift @@ -0,0 +1,45 @@ +import Foundation + +/// Compile-time constants for the MOSS-TTS-Nano backend. +/// +/// Geometry that is baked into the published CoreML graphs +/// (`FluidInference/moss-tts-nano-coreml`). Tokenizer-derived values (prompt +/// template ids, special token ids) are read from the repo's `config.json` +/// at load time instead — see `MossTtsNanoConfig`. +public enum MossTtsNanoConstants { + + // MARK: - Audio + + /// Codec output rate. Native 48 kHz stereo. + public static let sampleRate: Int = 48_000 + public static let channels: Int = 2 + /// One LM frame = 80 ms = 3840 samples per channel. + public static let samplesPerFramePerChannel: Int = 3840 + public static let frameDuration: Double = 0.08 + + // MARK: - Language model + + /// RVQ codebooks per frame (row width is `numCodebooks + 1`). + public static let numCodebooks: Int = 16 + public static let rowWidth: Int = 17 + public static let hiddenSize: Int = 768 + public static let audioCodebookSize: Int = 1024 + /// Fixed prompt capacity of the published prefill graph. + public static let prefillRows: Int = 512 + /// KV-cache capacity of the published step graph (prompt + generated frames). + public static let maxLen: Int = 1024 + + // MARK: - Generation defaults (upstream `inference()` defaults) + + public static let defaultTextTemperature: Float = 1.5 + public static let defaultAudioTemperature: Float = 1.7 + public static let defaultAudioTopP: Float = 0.8 + public static let defaultRepetitionPenalty: Float = 1.0 + public static let defaultMaxNewFrames: Int = 375 + /// Sentence-chunk token budget for voice-clone synthesis + /// (`DEFAULT_VOICE_CLONE_MAX_TEXT_TOKENS`). + public static let defaultMaxTextTokens: Int = 50 + /// Silence inserted between text chunks: short for ≤ 4-word chunks, long otherwise. + public static let interChunkPauseShort: Float = 0.40 + public static let interChunkPauseLong: Float = 0.24 +} diff --git a/Sources/FluidAudio/TTS/MossTtsNano/MossTtsNanoError.swift b/Sources/FluidAudio/TTS/MossTtsNano/MossTtsNanoError.swift new file mode 100644 index 000000000..c96a8881d --- /dev/null +++ b/Sources/FluidAudio/TTS/MossTtsNano/MossTtsNanoError.swift @@ -0,0 +1,50 @@ +import Foundation + +/// Errors surfaced by the MOSS-TTS-Nano backend. +public enum MossTtsNanoError: Error, LocalizedError, Sendable { + case notInitialized + case modelFileNotFound(String) + case corruptedModel(String, underlying: String) + case downloadFailed(String) + case configLoadFailed(String) + case tokenizerLoadFailed(String) + case voiceLoadFailed(path: String, underlying: String) + case invalidVoice(String) + case emptyText + case promptTooLong(rows: Int, capacity: Int) + case inferenceFailed(stage: String, underlying: String) + case invalidTensorShape(stage: String, expected: String, got: String) + case audioLoadFailed(path: String, underlying: String) + + public var errorDescription: String? { + switch self { + case .notInitialized: + return "MossTtsNano manager has not been initialized. Call initialize() first." + case .modelFileNotFound(let name): + return "MossTtsNano model file not found: \(name)" + case .corruptedModel(let name, let underlying): + return "MossTtsNano model appears corrupted: \(name) (\(underlying))" + case .downloadFailed(let message): + return "MossTtsNano download failed: \(message)" + case .configLoadFailed(let message): + return "MossTtsNano config.json load failed: \(message)" + case .tokenizerLoadFailed(let message): + return "MossTtsNano tokenizer.model load failed: \(message)" + case .voiceLoadFailed(let path, let underlying): + return "MossTtsNano voice load failed at \(path): \(underlying)" + case .invalidVoice(let message): + return "MossTtsNano invalid voice: \(message)" + case .emptyText: + return "MossTtsNano received empty text after normalization." + case .promptTooLong(let rows, let capacity): + return + "MossTtsNano prompt has \(rows) rows but the prefill graph holds \(capacity); use a shorter reference clip." + case .inferenceFailed(let stage, let underlying): + return "MossTtsNano \(stage) inference failed: \(underlying)" + case .invalidTensorShape(let stage, let expected, let got): + return "MossTtsNano \(stage) tensor shape mismatch: expected \(expected), got \(got)" + case .audioLoadFailed(let path, let underlying): + return "MossTtsNano could not load reference audio at \(path): \(underlying)" + } + } +} diff --git a/Sources/FluidAudio/TTS/MossTtsNano/MossTtsNanoManager.swift b/Sources/FluidAudio/TTS/MossTtsNano/MossTtsNanoManager.swift new file mode 100644 index 000000000..cc5e271d9 --- /dev/null +++ b/Sources/FluidAudio/TTS/MossTtsNano/MossTtsNanoManager.swift @@ -0,0 +1,188 @@ +@preconcurrency import CoreML +import Foundation + +/// Public API for MOSS-TTS-Nano: 0.1B multilingual streaming TTS with zero-shot +/// voice cloning, native 48 kHz stereo output. +/// +/// - Note: Beta — this is a beta model conversion; API, model artifacts, and accuracy may change. +/// +/// ```swift +/// let manager = try await MossTtsNanoManager.downloadAndCreate() +/// let voice = try await manager.loadVoice(.en2) +/// let audio = try await manager.synthesize(text: "Hello from the neural engine.", voice: voice) +/// // audio.left / audio.right are 48 kHz Float32; audio.mono for single-channel sinks +/// +/// for try await frame in try await manager.synthesizeStreaming(text: "…", voice: voice) { +/// player.schedule(left: frame.left, right: frame.right) // 80 ms per frame +/// } +/// ``` +/// +/// Voices are codec-token sequences of a reference clip: `loadVoice` fetches the +/// two published presets, `cloneVoice(audioURL:)` encodes any clip (≤ ~25 s) with +/// the fp32 codec encoder. Long text is chunked at sentence boundaries with the +/// upstream 50-token budget and joined with short pauses. +public actor MossTtsNanoManager { + + private let logger = AppLogger(category: "MossTtsNanoManager") + + private let directory: URL? + private let computeUnits: MLComputeUnits + private var store: MossTtsNanoModelStore? + private var synthesizer: MossTtsNanoSynthesizer? + + /// - Parameter computeUnits: unit preference for the Frame and CodecStep graphs + /// (Prefill/Step always run on CPU+GPU). `.cpuAndGPU` is the measured best. + public init(directory: URL? = nil, computeUnits: MLComputeUnits = .cpuAndGPU) { + self.directory = directory + self.computeUnits = computeUnits + } + + public var isAvailable: Bool { synthesizer != nil } + + public static func downloadAndCreate( + cacheDirectory: URL? = nil, + computeUnits: MLComputeUnits = .cpuAndGPU + ) async throws -> MossTtsNanoManager { + let manager = MossTtsNanoManager(directory: cacheDirectory, computeUnits: computeUnits) + try await manager.initialize() + return manager + } + + /// Download (if missing) and load the streaming-path models. + public func initialize() async throws { + if synthesizer != nil { return } + let store = MossTtsNanoModelStore(directory: directory, computeUnits: computeUnits) + try await store.loadIfNeeded() + synthesizer = MossTtsNanoSynthesizer( + prefill: try await store.prefill(), + step: try await store.step(), + frame: try await store.frame(), + codecStep: try await store.codecStep(), + config: try await store.config(), + tokenizer: try await store.tokenizer()) + self.store = store + logger.info("MOSS-TTS-Nano ready") + } + + public func cleanup() async { + if let store { await store.unload() } + store = nil + synthesizer = nil + } + + // MARK: - Voices + + /// Fetch and decode a published preset voice. + public func loadVoice(_ voice: MossTtsNanoBuiltInVoice) async throws -> MossTtsNanoVoice { + try await MossTtsNanoResourceDownloader.loadVoice(voice, directory: directory) + } + + /// Encode a reference clip into a reusable voice. Any AVFoundation-readable file; + /// resampled to 48 kHz stereo internally. Downloads the fp32 codec encoder on first use. + public func cloneVoice(audioURL: URL, name: String? = nil) async throws -> MossTtsNanoVoice { + guard let store, let synthesizer else { throw MossTtsNanoError.notInitialized } + let encoder = try await store.encoder() + let (left, right) = try MossTtsNanoAudioLoader.loadStereo48k(url: audioURL) + let voice = try MossTtsNanoSynthesizer.encodeVoice( + encoder: encoder, left: left, right: right, + name: name ?? audioURL.deletingPathExtension().lastPathComponent, + nVq: synthesizer.config.model.nVq, + samplesPerFrame: synthesizer.config.model.samplesPerFramePerChannel) + let overhead = synthesizer.builder.voiceCloneOverheadRows(voice: voice) + let capacity = synthesizer.config.coreml.prefillRows + guard overhead < capacity - 8 else { + throw MossTtsNanoError.invalidVoice( + "reference clip is \(voice.frames) frames (\(Double(voice.frames) * MossTtsNanoConstants.frameDuration) s); " + + "it leaves no room for text in the \(capacity)-row prompt. Use a clip under ~25 s.") + } + return voice + } + + // MARK: - Synthesis + + /// Tokenize text exactly as the model prompt does (SentencePiece BPE ids). + public func tokenize(_ text: String) throws -> [Int] { + guard let synthesizer else { throw MossTtsNanoError.notInitialized } + return synthesizer.tokenizer.encode(text) + } + + /// Synthesize a full utterance (48 kHz stereo). + public func synthesize( + text: String, + voice: MossTtsNanoVoice, + options: MossTtsNanoSamplingOptions = .default + ) async throws -> MossTtsNanoAudio { + var left: [Float] = [] + var right: [Float] = [] + let stream = try synthesizeStreaming(text: text, voice: voice, options: options) + for try await frame in stream { + left.append(contentsOf: frame.left) + right.append(contentsOf: frame.right) + } + return MossTtsNanoAudio(left: left, right: right) + } + + /// Synthesize as a stream of 80 ms stereo frames, yielded as soon as each is decoded. + /// Pause frames (`isPause == true`) separate text chunks. + public func synthesizeStreaming( + text: String, + voice: MossTtsNanoVoice, + options: MossTtsNanoSamplingOptions = .default + ) throws -> AsyncThrowingStream { + guard let synthesizer else { throw MossTtsNanoError.notInitialized } + try voice.validate() + let chunks = try synthesizer.planChunks(text: text, voice: voice, options: options) + logger.info("MOSS-TTS-Nano synthesizing \(chunks.count) chunk(s), voice=\(voice.name)") + return AsyncThrowingStream { continuation in + let task = Task { + do { + try await self.run(chunks: chunks, voice: voice, options: options, continuation: continuation) + continuation.finish() + } catch { + continuation.finish(throwing: error) + } + } + continuation.onTermination = { _ in task.cancel() } + } + } + + private func run( + chunks: [String], + voice: MossTtsNanoVoice, + options: MossTtsNanoSamplingOptions, + continuation: AsyncThrowingStream.Continuation + ) async throws { + guard let synthesizer else { throw MossTtsNanoError.notInitialized } + var rng = options.seed.map { MossTtsNanoRandom(seed: $0) } ?? MossTtsNanoRandom() + let start = Date() + var totalFrames = 0 + for (index, chunk) in chunks.enumerated() { + try Task.checkCancellation() + let produced = try synthesizer.generateChunk( + text: chunk, voice: voice, options: options, rng: &rng, + chunkIndex: index, chunkCount: chunks.count + ) { frame in + try Task.checkCancellation() + continuation.yield(frame) + } + totalFrames += produced + if index < chunks.count - 1 { + let pause = MossTtsNanoTextChunker.pauseSeconds(after: chunk) + let samples = Int((pause * Float(MossTtsNanoConstants.sampleRate)).rounded()) + if samples > 0 { + let silence = [Float](repeating: 0, count: samples) + continuation.yield( + MossTtsNanoAudioFrame( + left: silence, right: silence, frameIndex: produced, + chunkIndex: index, chunkCount: chunks.count, isPause: true)) + } + } + } + let elapsed = Date().timeIntervalSince(start) + let audioSeconds = Double(totalFrames) * MossTtsNanoConstants.frameDuration + logger.info( + "MOSS-TTS-Nano generated \(totalFrames) frames (\(String(format: "%.2f", audioSeconds)) s) in " + + "\(String(format: "%.2f", elapsed)) s (RTFx \(String(format: "%.2f", elapsed > 0 ? audioSeconds / elapsed : 0)))" + ) + } +} diff --git a/Sources/FluidAudio/TTS/MossTtsNano/MossTtsNanoTypes.swift b/Sources/FluidAudio/TTS/MossTtsNano/MossTtsNanoTypes.swift new file mode 100644 index 000000000..e3a105e8f --- /dev/null +++ b/Sources/FluidAudio/TTS/MossTtsNano/MossTtsNanoTypes.swift @@ -0,0 +1,229 @@ +import Foundation + +/// Mirror of `config.json` in `FluidInference/moss-tts-nano-coreml`. +/// +/// Carries everything the Swift host needs that is derived from the upstream +/// tokenizer / model config rather than baked into the CoreML graphs: special +/// token ids, pre-tokenized prompt template segments, sampling defaults. +public struct MossTtsNanoConfig: Codable, Sendable { + + public struct Model: Codable, Sendable { + public let nVq: Int + public let rowWidth: Int + public let hiddenSize: Int + public let audioCodebookSize: Int + public let sampleRate: Int + public let channels: Int + public let samplesPerFramePerChannel: Int + } + + public struct CoreML: Codable, Sendable { + public let prefillRows: Int + public let maxLen: Int + } + + public struct Tokens: Codable, Sendable { + public let pad: Int + public let imStart: Int + public let imEnd: Int + public let audioStart: Int + public let audioEnd: Int + public let audioUserSlot: Int + public let audioAssistantSlot: Int + public let audioPad: Int + } + + /// Pre-tokenized template segments (upstream `prompting.py`). + public struct Prompt: Codable, Sendable { + /// ` user … Reference(s): ` + public let voiceClonePrefix: [Int] + /// ` … Text:` — sits between the reference rows and the text ids. + public let voiceCloneAfterReference: [Int] + /// ` assistant ` + public let assistantSuffix: [Int] + /// Plain (no reference) prompt prefix; kept for completeness. + public let plainPrefix: [Int] + } + + public struct SamplingDefaults: Codable, Sendable { + public let textTemperature: Float + public let textTopK: Int + public let textTopP: Float + public let audioTemperature: Float + public let audioTopK: Int + public let audioTopP: Float + public let audioRepetitionPenalty: Float + public let maxNewFrames: Int + } + + public struct Text: Codable, Sendable { + public let voiceCloneMaxTextTokens: Int + } + + public let model: Model + public let coreml: CoreML + public let tokens: Tokens + public let prompt: Prompt + public let samplingDefaults: SamplingDefaults + public let text: Text + + public static func load(from url: URL) throws -> MossTtsNanoConfig { + let data = try Data(contentsOf: url) + let decoder = JSONDecoder() + decoder.keyDecodingStrategy = .convertFromSnakeCase + return try decoder.decode(MossTtsNanoConfig.self, from: data) + } +} + +/// A voice-clone reference: the codec tokens of a short speech clip. +/// +/// `codes` is `[frames][16]` (12.5 Hz frames × RVQ codebooks). Built-in voices +/// ship as `voices/.json` in the HF repo; custom voices come from +/// `MossTtsNanoManager.cloneVoice(audioURL:)`. +public struct MossTtsNanoVoice: Codable, Sendable, Equatable { + public let name: String + public let sampleRate: Int + public let frames: Int + public let codes: [[Int32]] + + public init(name: String, codes: [[Int32]], sampleRate: Int = MossTtsNanoConstants.sampleRate) { + self.name = name + self.sampleRate = sampleRate + self.frames = codes.count + self.codes = codes + } + + public static func load(from url: URL) throws -> MossTtsNanoVoice { + do { + let data = try Data(contentsOf: url) + let decoder = JSONDecoder() + decoder.keyDecodingStrategy = .convertFromSnakeCase + let voice = try decoder.decode(MossTtsNanoVoice.self, from: data) + try voice.validate() + return voice + } catch let error as MossTtsNanoError { + throw error + } catch { + throw MossTtsNanoError.voiceLoadFailed(path: url.path, underlying: "\(error)") + } + } + + public func save(to url: URL) throws { + let encoder = JSONEncoder() + encoder.keyEncodingStrategy = .convertToSnakeCase + try encoder.encode(self).write(to: url) + } + + func validate() throws { + guard !codes.isEmpty else { throw MossTtsNanoError.invalidVoice("no frames") } + for (i, row) in codes.enumerated() { + guard row.count == MossTtsNanoConstants.numCodebooks else { + throw MossTtsNanoError.invalidVoice( + "frame \(i) has \(row.count) codes, expected \(MossTtsNanoConstants.numCodebooks)") + } + guard row.allSatisfy({ $0 >= 0 && $0 < Int32(MossTtsNanoConstants.audioCodebookSize) }) else { + throw MossTtsNanoError.invalidVoice("frame \(i) has a code outside 0..<1024") + } + } + } +} + +/// Voices published alongside the models (`voices/.json`). +public enum MossTtsNanoBuiltInVoice: String, CaseIterable, Sendable { + /// English female speaker (upstream demo clip `en_2.wav`). + case en2 = "en_2" + /// Mandarin speaker (upstream demo clip `zh_1.wav`). + case zh1 = "zh_1" + + public static let `default`: MossTtsNanoBuiltInVoice = .en2 + + public var fileName: String { "\(MossTtsNanoBuiltInVoiceFiles.subdirectory)/\(rawValue).json" } + + public init?(name: String) { + self.init(rawValue: name.lowercased().replacingOccurrences(of: "-", with: "_")) + } +} + +enum MossTtsNanoBuiltInVoiceFiles { + static let subdirectory = "voices" +} + +/// Sampling controls. Defaults mirror upstream `inference()`. +public struct MossTtsNanoSamplingOptions: Sendable, Equatable { + public var textTemperature: Float + public var audioTemperature: Float + public var audioTopP: Float + public var repetitionPenalty: Float + /// Argmax decoding. Deterministic, but upstream greedy decoding never emits the + /// stop token and runs to `maxNewFrames`; it exists for parity testing only. + public var greedy: Bool + public var maxNewFrames: Int + /// Seed for the host-side uniform draws that feed the in-graph sampler. `nil` + /// draws from the system RNG. + public var seed: UInt64? + /// Token budget per text chunk; `0` disables chunking. + public var maxTextTokens: Int + + public init( + textTemperature: Float = MossTtsNanoConstants.defaultTextTemperature, + audioTemperature: Float = MossTtsNanoConstants.defaultAudioTemperature, + audioTopP: Float = MossTtsNanoConstants.defaultAudioTopP, + repetitionPenalty: Float = MossTtsNanoConstants.defaultRepetitionPenalty, + greedy: Bool = false, + maxNewFrames: Int = MossTtsNanoConstants.defaultMaxNewFrames, + seed: UInt64? = nil, + maxTextTokens: Int = MossTtsNanoConstants.defaultMaxTextTokens + ) { + self.textTemperature = textTemperature + self.audioTemperature = audioTemperature + self.audioTopP = audioTopP + self.repetitionPenalty = repetitionPenalty + self.greedy = greedy + self.maxNewFrames = maxNewFrames + self.seed = seed + self.maxTextTokens = maxTextTokens + } + + public static let `default` = MossTtsNanoSamplingOptions() +} + +/// One decoded 80 ms frame (3840 samples per channel at 48 kHz). +public struct MossTtsNanoAudioFrame: Sendable { + public let left: [Float] + public let right: [Float] + /// Zero-based frame index within the current text chunk. + public let frameIndex: Int + /// Zero-based text chunk index and total chunk count for the utterance. + public let chunkIndex: Int + public let chunkCount: Int + /// Silence frame inserted between chunks (not model output). + public let isPause: Bool + public var sampleRate: Int { MossTtsNanoConstants.sampleRate } +} + +/// A complete stereo utterance. +public struct MossTtsNanoAudio: Sendable { + public let left: [Float] + public let right: [Float] + public var sampleRate: Int { MossTtsNanoConstants.sampleRate } + public var duration: Double { Double(left.count) / Double(sampleRate) } + + /// Mid channel `(L + R) / 2`. + public var mono: [Float] { + var out = [Float](repeating: 0, count: left.count) + for i in 0..) +/// [audio_user_slot, codes…] × reference frames +/// voice_clone_after_reference ( … Text:) +/// text ids +/// assistant_suffix (… assistant ) +/// ``` +struct MossTtsNanoPromptBuilder: Sendable { + let config: MossTtsNanoConfig + + var rowWidth: Int { config.model.rowWidth } + + func textRow(_ tokenId: Int) -> [Int32] { + var row = [Int32](repeating: Int32(config.tokens.audioPad), count: rowWidth) + row[0] = Int32(tokenId) + return row + } + + func generationRow(codes: [Int32]) -> [Int32] { + var row = [Int32](repeating: Int32(config.tokens.audioPad), count: rowWidth) + row[0] = Int32(config.tokens.audioAssistantSlot) + for (i, c) in codes.prefix(rowWidth - 1).enumerated() { + row[i + 1] = c + } + return row + } + + func voiceCloneRows(textIds: [Int], voice: MossTtsNanoVoice) -> [[Int32]] { + var rows: [[Int32]] = [] + rows.reserveCapacity( + config.prompt.voiceClonePrefix.count + voice.frames + config.prompt.voiceCloneAfterReference.count + + textIds.count + config.prompt.assistantSuffix.count) + for id in config.prompt.voiceClonePrefix { rows.append(textRow(id)) } + for codes in voice.codes { + var row = [Int32](repeating: Int32(config.tokens.audioPad), count: rowWidth) + row[0] = Int32(config.tokens.audioUserSlot) + for (i, c) in codes.prefix(rowWidth - 1).enumerated() { + row[i + 1] = c + } + rows.append(row) + } + for id in config.prompt.voiceCloneAfterReference { rows.append(textRow(id)) } + for id in textIds { rows.append(textRow(id)) } + for id in config.prompt.assistantSuffix { rows.append(textRow(id)) } + return rows + } + + /// Rows that are not text: fixed template overhead for a given voice. + func voiceCloneOverheadRows(voice: MossTtsNanoVoice) -> Int { + config.prompt.voiceClonePrefix.count + voice.frames + config.prompt.voiceCloneAfterReference.count + + config.prompt.assistantSuffix.count + } +} diff --git a/Sources/FluidAudio/TTS/MossTtsNano/Pipeline/MossTtsNanoSynthesizer.swift b/Sources/FluidAudio/TTS/MossTtsNano/Pipeline/MossTtsNanoSynthesizer.swift new file mode 100644 index 000000000..dc1bf9ff4 --- /dev/null +++ b/Sources/FluidAudio/TTS/MossTtsNano/Pipeline/MossTtsNanoSynthesizer.swift @@ -0,0 +1,228 @@ +@preconcurrency import CoreML +import Foundation + +/// Drives the MOSS-TTS-Nano CoreML graphs frame by frame. +/// +/// Per text chunk: +/// 1. `Prefill(input_ids [1,512,17], input_len)` → `hidden`, `kv_k`, `kv_v` +/// 2. loop: `Frame(hidden, randoms, sampling params, seen, greedy)` → +/// `should_continue`, 16 codes → `CodecStep(codes, frame_index, caches…)` → +/// 3840 stereo samples → `Step(row, kv, cur_len)` → next `hidden`, kv. +/// +/// Sampling happens inside the Frame graph (top-k / top-p / temperature / +/// repetition penalty via inverse-CDF); the host only supplies uniform randoms. +final class MossTtsNanoSynthesizer { + + private let logger = AppLogger(category: "MossTtsNanoSynthesizer") + + private let prefill: MLModel + private let step: MLModel + private let frame: MLModel + private let codecStep: MLModel + let config: MossTtsNanoConfig + let tokenizer: MossTtsNanoTokenizer + let builder: MossTtsNanoPromptBuilder + private let codecCacheShapes: [(name: String, shape: [NSNumber])] + + init( + prefill: MLModel, step: MLModel, frame: MLModel, codecStep: MLModel, + config: MossTtsNanoConfig, tokenizer: MossTtsNanoTokenizer + ) { + self.prefill = prefill + self.step = step + self.frame = frame + self.codecStep = codecStep + self.config = config + self.tokenizer = tokenizer + self.builder = MossTtsNanoPromptBuilder(config: config) + self.codecCacheShapes = codecStep.modelDescription.inputDescriptionsByName + .filter { $0.key != "codes" && $0.key != "frame_index" } + .compactMap { name, desc in + desc.multiArrayConstraint.map { (name: name, shape: $0.shape) } + } + .sorted { $0.name < $1.name } + } + + // MARK: - Planning + + /// Split `text` into chunks that fit the prefill graph next to `voice`. + func planChunks(text: String, voice: MossTtsNanoVoice, options: MossTtsNanoSamplingOptions) throws -> [String] { + let overhead = builder.voiceCloneOverheadRows(voice: voice) + let capacity = config.coreml.prefillRows - overhead + guard capacity > 0 else { + throw MossTtsNanoError.promptTooLong(rows: overhead, capacity: config.coreml.prefillRows) + } + var budget = options.maxTextTokens > 0 ? options.maxTextTokens : capacity + budget = min(budget, capacity) + let chunks = try MossTtsNanoTextChunker.chunk(text, maxTokens: budget) { [tokenizer] in + tokenizer.encode($0).count + } + for chunk in chunks { + let rows = overhead + tokenizer.encode(chunk).count + if rows > config.coreml.prefillRows { + throw MossTtsNanoError.promptTooLong(rows: rows, capacity: config.coreml.prefillRows) + } + } + return chunks + } + + // MARK: - Generation + + /// Generate one chunk, calling `emit` per decoded 80 ms frame. Returns the frame count. + func generateChunk( + text: String, + voice: MossTtsNanoVoice, + options: MossTtsNanoSamplingOptions, + rng: inout MossTtsNanoRandom, + chunkIndex: Int, + chunkCount: Int, + emit: (MossTtsNanoAudioFrame) throws -> Void + ) throws -> Int { + let nVq = config.model.nVq + let rowWidth = config.model.rowWidth + let prefillRows = config.coreml.prefillRows + let textIds = tokenizer.encode(text) + guard !textIds.isEmpty else { throw MossTtsNanoError.emptyText } + let rows = builder.voiceCloneRows(textIds: textIds, voice: voice) + guard rows.count <= prefillRows else { + throw MossTtsNanoError.promptTooLong(rows: rows.count, capacity: prefillRows) + } + + // --- Prefill --- // + var padded = [Int32](repeating: Int32(config.tokens.audioPad), count: prefillRows * rowWidth) + for r in 0..= rows.count { + padded[r * rowWidth] = Int32(config.tokens.pad) + } + for (r, row) in rows.enumerated() { + for (c, v) in row.enumerated() { padded[r * rowWidth + c] = v } + } + let prefillOut = try MossTtsNanoTensor.predict( + prefill, + [ + "input_ids": try MossTtsNanoTensor.int32(padded, shape: [1, prefillRows, rowWidth], stage: "prefill"), + "input_len": try MossTtsNanoTensor.int32([Int32(rows.count)], shape: [1], stage: "prefill"), + ], stage: "prefill") + var hidden = try MossTtsNanoTensor.output(prefillOut, "hidden", stage: "prefill") + var kvK = try MossTtsNanoTensor.output(prefillOut, "kv_k", stage: "prefill") + var kvV = try MossTtsNanoTensor.output(prefillOut, "kv_v", stage: "prefill") + + // --- Per-frame state --- // + var caches: [String: MLMultiArray] = [:] + for (name, shape) in codecCacheShapes { + caches[name] = try MossTtsNanoTensor.zeros(shape: shape) + } + let seen = try MossTtsNanoTensor.zeros(shape: [1, nVq, config.model.audioCodebookSize].map(NSNumber.init)) + let seenPtr = seen.dataPointer.bindMemory(to: Float.self, capacity: nVq * config.model.audioCodebookSize) + let textTemperature = try MossTtsNanoTensor.float32([options.textTemperature], shape: [1], stage: "frame") + let audioTemperature = try MossTtsNanoTensor.float32([options.audioTemperature], shape: [1], stage: "frame") + let audioTopP = try MossTtsNanoTensor.float32([options.audioTopP], shape: [1], stage: "frame") + let repetitionPenalty = try MossTtsNanoTensor.float32([options.repetitionPenalty], shape: [1], stage: "frame") + let greedy = try MossTtsNanoTensor.float32([options.greedy ? 1 : 0], shape: [1], stage: "frame") + + let maxFrames = min(options.maxNewFrames, config.coreml.maxLen - rows.count - 1) + var curLen = rows.count + var produced = 0 + let samplesPerChannel = config.model.samplesPerFramePerChannel + + for t in 0..= 0 && Int(code) < config.model.audioCodebookSize { + seenPtr[c * config.model.audioCodebookSize + Int(code)] = 1 + } + + // --- Codec step: 16 codes → 80 ms stereo --- // + var codecInputs: [String: MLMultiArray] = [ + "codes": try MossTtsNanoTensor.int32(codes, shape: [nVq, 1, 1], stage: "codec_step"), + "frame_index": try MossTtsNanoTensor.int32([Int32(t)], shape: [1], stage: "codec_step"), + ] + for (name, cache) in caches { codecInputs[name] = cache } + let codecOut = try MossTtsNanoTensor.predict(codecStep, codecInputs, stage: "codec_step") + let audio = MossTtsNanoTensor.floats(try MossTtsNanoTensor.output(codecOut, "audio", stage: "codec_step")) + guard audio.count == 2 * samplesPerChannel else { + throw MossTtsNanoError.invalidTensorShape( + stage: "codec_step", expected: "\(2 * samplesPerChannel)", got: "\(audio.count)") + } + for (name, _) in codecCacheShapes { + caches[name] = try MossTtsNanoTensor.output(codecOut, "\(name)_out", stage: "codec_step") + } + try emit( + MossTtsNanoAudioFrame( + left: Array(audio[0..