From a9ebc176cb3ee9b503790acf5dcd66609712b34b Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Wed, 19 Aug 2026 13:33:43 -0400 Subject: [PATCH 1/2] fix(tts/pockettts): per-language mimi encoders for non-English voice cloning (#793) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause of the residual flaky/garbled non-English cloning: every pocket-tts language pack ships its OWN mimi codec weights (all 87 mimi state-dict keys differ across languages, including the 43 encoder-side keys). The shared root mimi_encoderv2 was traced from the English model, so live-clone conditioning for non-English packs was built from the wrong codec's latents. The #797 reprojection corrected the projection but not the latents — Swift conditioning differed from the reference _encode_audio output by ~140% relative, and injecting it into the PyTorch reference reproduced the exact failure signature (no-EOS runaway noise at low temperature, wrong-voice timbre), while native conditioning was 16/16 clean with the same 5s reference clip. Fix: prefer a pack-local per-language encoder (v2.1//mimi_encoderv3.mlmodelc, that pack's mimi + speaker_proj baked in, published on FluidInference/pocket-tts-coreml) fetched lazily on first clone. Its output is already in the pack's own conditioning space, so no host-side reprojection is applied. The shared root encoder + reprojection remains the fallback (English, stale caches, offline). Swift e2e with the issue #793 reporter's reference on spanish_24l goes from 7/16 intelligible (4 pure-noise runaways at temperature 0.15) to 15/16 — on par with the PyTorch reference under the same seeds. 6-layer spanish cloning goes from always-garbled to 3/4 (upstream-level), so the 6L warn-and-steer now only fires on the fallback path. Also adds a --temperature flag to the tts CLI (PocketTTS path) to make reliability matrices reproducible. --- Sources/FluidAudio/ModelNames.swift | 8 ++ .../Assets/PocketTtsResourceDownloader.swift | 52 +++++++++++++ .../Pipeline/PocketTtsModelStore.swift | 76 ++++++++++++++----- .../FluidAudioCLI/Commands/TTSCommand.swift | 32 +++++--- .../PocketTtsDownloadFilterTests.swift | 10 +++ 5 files changed, 148 insertions(+), 30 deletions(-) diff --git a/Sources/FluidAudio/ModelNames.swift b/Sources/FluidAudio/ModelNames.swift index edfa91bdb..5abe97f84 100644 --- a/Sources/FluidAudio/ModelNames.swift +++ b/Sources/FluidAudio/ModelNames.swift @@ -998,6 +998,13 @@ public enum ModelNames { public static let pocketState = "pocket_state" public static let mimiDecoder = "mimi_decoder" public static let mimiEncoder = "mimi_encoderv2" + /// Per-language voice-cloning encoder, stored inside each pack's + /// directory (`v2.1//mimi_encoderv3.mlmodelc`). Every language + /// pack ships its own mimi weights, so cloned-voice conditioning must + /// be encoded with the pack's own codec + speaker projection (baked + /// in at conversion). The shared root `mimi_encoderv2` (English mimi) + /// remains the fallback for English and stale caches (#793). + public static let mimiEncoderV3 = "mimi_encoderv3" /// Function names inside the `pocket_state` multifunction package. public enum StateFunction { @@ -1017,6 +1024,7 @@ public enum ModelNames { public static let pocketStateFile = pocketState + ".mlmodelc" public static let mimiDecoderFile = mimiDecoder + ".mlmodelc" public static let mimiEncoderFile = mimiEncoder + ".mlmodelc" + public static let mimiEncoderV3File = mimiEncoderV3 + ".mlmodelc" /// Directory containing binary constants, tokenizer, and voice data. public static let constantsBinDir = "constants_bin" diff --git a/Sources/FluidAudio/TTS/PocketTTS/Assets/PocketTtsResourceDownloader.swift b/Sources/FluidAudio/TTS/PocketTTS/Assets/PocketTtsResourceDownloader.swift index ffba2400b..24f9de8f7 100644 --- a/Sources/FluidAudio/TTS/PocketTTS/Assets/PocketTtsResourceDownloader.swift +++ b/Sources/FluidAudio/TTS/PocketTTS/Assets/PocketTtsResourceDownloader.swift @@ -349,6 +349,58 @@ public enum PocketTtsResourceDownloader { return encoderPath } + /// Ensure the pack-local per-language voice-cloning encoder + /// (`v2.1//mimi_encoderv3.mlmodelc`) is available and return its URL. + /// + /// Every language pack ships its own mimi codec weights, so cloned-voice + /// conditioning must be encoded with the pack's own encoder (its speaker + /// projection is baked in — no host-side reprojection). Best-effort: + /// returns `nil` when the encoder can't be fetched (offline, or not yet + /// published for the pack); callers then fall back to the shared root + /// encoder + reprojection path (#793). + public static func ensurePackMimiEncoder( + language: PocketTtsLanguage, directory: URL? = nil + ) async -> URL? { + do { + let targetDir = try directory ?? cacheDirectory() + let modelsDirectory = targetDir.appendingPathComponent( + PocketTtsConstants.defaultModelsSubdirectory) + let repoDir = modelsDirectory.appendingPathComponent(Repo.pocketTts.folderName) + let encoderSubpath = + "\(language.repoSubdirectory)/\(ModelNames.PocketTTS.mimiEncoderV3File)" + let encoderPath = repoDir.appendingPathComponent(encoderSubpath) + + if FileManager.default.fileExists(atPath: encoderPath.path) { + return encoderPath + } + + try FileManager.default.createDirectory( + at: repoDir, withIntermediateDirectories: true) + + logger.info( + "Downloading per-language Mimi encoder for \(language.rawValue) voice cloning...") + try await ModelHub.download( + .pocketTts, + subdirectory: encoderSubpath, + to: repoDir + ) + + guard FileManager.default.fileExists(atPath: encoderPath.path) else { + logger.warning( + "Per-language Mimi encoder unavailable for \(language.rawValue); " + + "falling back to the shared encoder + reprojection (#793).") + return nil + } + return encoderPath + } catch { + logger.warning( + "Failed to fetch per-language Mimi encoder for \(language.rawValue): " + + "\(error.localizedDescription). Falling back to the shared encoder + " + + "reprojection (#793).") + return nil + } + } + /// Ensure voice conditioning data for the given language is available, /// downloading from HuggingFace if missing. /// diff --git a/Sources/FluidAudio/TTS/PocketTTS/Pipeline/PocketTtsModelStore.swift b/Sources/FluidAudio/TTS/PocketTTS/Pipeline/PocketTtsModelStore.swift index a70c135cd..67775ac52 100644 --- a/Sources/FluidAudio/TTS/PocketTTS/Pipeline/PocketTtsModelStore.swift +++ b/Sources/FluidAudio/TTS/PocketTTS/Pipeline/PocketTtsModelStore.swift @@ -19,6 +19,10 @@ public actor PocketTtsModelStore { private var flowDecoderModel: MLModel? private var mimiDecoderModel: MLModel? private var mimiEncoderModel: MLModel? + /// True when `mimiEncoderModel` is the pack-local per-language encoder + /// (`mimi_encoderv3.mlmodelc`), whose output is already in the pack's own + /// conditioning space — no host-side reprojection needed (#793). + private var mimiEncoderIsPackLocal = false private var speakerProjectionCache: PocketTtsVoiceCloner.SpeakerProjection? /// `.aneState` only: the prefill/generate function instances loaded from /// the `pocket_state.mlmodelc` multifunction package (mobius Trial 23). @@ -379,21 +383,38 @@ public actor PocketTtsModelStore { /// Load the Mimi encoder model for voice cloning (lazy, on-demand). /// - /// Downloads the model from HuggingFace if not already cached. The Mimi - /// encoder is language-agnostic and lives at the repo root, shared - /// across all language packs. + /// Downloads the model from HuggingFace if not already cached. Every + /// language pack ships its own mimi codec weights, so non-English packs + /// prefer the pack-local `mimi_encoderv3.mlmodelc` (that pack's mimi + + /// speaker projection baked in). The shared root `mimi_encoderv2` (English + /// mimi) is English's encoder and the fallback when the pack-local one is + /// unavailable — cloning then goes through the legacy reprojection path, + /// which places the voice only approximately (#793). public func loadMimiEncoderIfNeeded() async throws { guard mimiEncoderModel == nil else { return } - // Ensure the mimi_encoder is downloaded (downloads if needed) - let modelURL = try await PocketTtsResourceDownloader.ensureMimiEncoder(directory: directory) - let config = MLModelConfiguration() config.computeUnits = .cpuAndGPU + if language != .english, + let packURL = await PocketTtsResourceDownloader.ensurePackMimiEncoder( + language: language, directory: directory) + { + logger.info("Loading pack-local Mimi encoder (\(self.language.rawValue))...") + let loadStart = Date() + mimiEncoderModel = try MLModel(contentsOf: packURL, configuration: config) + mimiEncoderIsPackLocal = true + let elapsed = Date().timeIntervalSince(loadStart) + logger.info("Mimi encoder loaded in \(String(format: "%.2f", elapsed))s") + return + } + + let modelURL = try await PocketTtsResourceDownloader.ensureMimiEncoder(directory: directory) + logger.info("Loading Mimi encoder for voice cloning...") let loadStart = Date() mimiEncoderModel = try MLModel(contentsOf: modelURL, configuration: config) + mimiEncoderIsPackLocal = false let elapsed = Date().timeIntervalSince(loadStart) logger.info("Mimi encoder loaded in \(String(format: "%.2f", elapsed))s") } @@ -408,11 +429,16 @@ public actor PocketTtsModelStore { return model } - /// Check if the Mimi encoder model is available. + /// Check if a Mimi encoder model is available (pack-local per-language + /// encoder inside the language root, or the shared root encoder). public func isMimiEncoderAvailable() -> Bool { - // The Mimi encoder lives at the repo root, two levels above any - // `v2//` language root. guard let langRoot = languageRootDirectory else { return false } + let packURL = langRoot.appendingPathComponent(ModelNames.PocketTTS.mimiEncoderV3File) + if FileManager.default.fileExists(atPath: packURL.path) { + return true + } + // The shared root encoder lives two levels above any `v2.1//` + // language root. let repoRoot = langRoot.deletingLastPathComponent().deletingLastPathComponent() let modelURL = repoRoot.appendingPathComponent(ModelNames.PocketTTS.mimiEncoderFile) return FileManager.default.fileExists(atPath: modelURL.path) @@ -422,14 +448,22 @@ public actor PocketTtsModelStore { public func cloneVoice(from audioURL: URL) throws -> PocketTtsVoiceData { let encoder = try mimiEncoder() return try PocketTtsVoiceCloner.cloneVoice( - from: audioURL, using: encoder, projection: loadSpeakerProjection()) + from: audioURL, using: encoder, projection: cloneProjection()) } /// Clone a voice from audio samples within the actor's isolation context. public func cloneVoice(from samples: [Float]) throws -> PocketTtsVoiceData { let encoder = try mimiEncoder() return try PocketTtsVoiceCloner.cloneVoice( - from: samples, using: encoder, projection: loadSpeakerProjection()) + from: samples, using: encoder, projection: cloneProjection()) + } + + /// Projection to apply to freshly-encoded clone conditioning. The + /// pack-local per-language encoder already emits conditioning in the + /// pack's own space, so no reprojection is applied; the legacy shared + /// (English) encoder needs the #793 reprojection assets. + private func cloneProjection() -> PocketTtsVoiceCloner.SpeakerProjection? { + mimiEncoderIsPackLocal ? nil : loadSpeakerProjection() } /// Load (and cache) the per-language speaker projection used to re-project @@ -460,17 +494,19 @@ public actor PocketTtsModelStore { guard let pinv = loadFloatBin(pinvURL, expectedCount: expected), let proj = loadFloatBin(projURL, expectedCount: expected) else { - // 6-layer non-English packs can't clone reliably even with the - // projection — the upstream pocket-tts flow LM early-EOSes on the - // encoded voice conditioning (confirmed against the PyTorch - // reference, #793). Steer callers to the 24-layer variant, which - // works. Other missing-asset cases are a transient fetch gap. + // This fallback only runs when the pack-local per-language + // encoder (`mimi_encoderv3`) is unavailable. The legacy shared + // (English-mimi) encoder + reprojection cannot clone reliably for + // 6-layer non-English packs — the flow LM early-EOSes on the + // wrong-codec conditioning (#793). Other missing-asset cases are + // a transient fetch gap. if language.transformerLayers == 6 && language != .english { logger.warning( - "PocketTTS live voice cloning is not supported for the \(self.language.rawValue) " - + "(6-layer) pack — the upstream model produces garbled/truncated output " - + "for cloned voices. Use the 24-layer variant (\(self.language.rawValue)_24l) " - + "for voice cloning (#793).") + "PocketTTS live voice cloning for the \(self.language.rawValue) (6-layer) pack " + + "requires the pack-local \(ModelNames.PocketTTS.mimiEncoderV3File), which " + + "is unavailable — the shared-encoder fallback produces garbled/truncated " + + "output for this pack. Retry once the encoder can be downloaded, or use " + + "the 24-layer variant (\(self.language.rawValue)_24l) (#793).") } else if language != .english { logger.warning( "PocketTTS speaker projection assets missing for \(self.language.rawValue) " diff --git a/Sources/FluidAudioCLI/Commands/TTSCommand.swift b/Sources/FluidAudioCLI/Commands/TTSCommand.swift index 0306d3088..9b6596106 100644 --- a/Sources/FluidAudioCLI/Commands/TTSCommand.swift +++ b/Sources/FluidAudioCLI/Commands/TTSCommand.swift @@ -80,6 +80,7 @@ public struct TTS { var saveVoicePath: String? = nil var pocketLanguage: PocketTtsLanguage = .english var pocketPlacement: PocketTtsModelPlacement = .gpu + var pocketTemperature: Float = PocketTtsConstants.temperature // PocketTTS deterministic-seed mode (uses session API for fixed RNG). var pocketSeed: UInt64? = nil // StyleTTS2 zero-shot args. @@ -228,6 +229,11 @@ public struct TTS { luxttsPromptText = arguments[i + 1] i += 1 } + case "--temperature": + if i + 1 < arguments.count, let v = Float(arguments[i + 1]) { + pocketTemperature = v + i += 1 + } case "--silence": if i + 1 < arguments.count, let v = Float(arguments[i + 1]) { supertonicSilence = v @@ -345,7 +351,7 @@ public struct TTS { metricsPath: metricsPath, cloneVoicePath: cloneVoicePath, voiceFilePath: voiceFilePath, saveVoicePath: saveVoicePath, language: pocketLanguage, seed: pocketSeed, - placement: pocketPlacement) + placement: pocketPlacement, temperature: pocketTemperature) case .kokoroAne: await runKokoroAne( text: text, output: output, voice: voice, metricsPath: metricsPath, @@ -638,11 +644,13 @@ public struct TTS { voice: String, voiceData: PocketTtsVoiceData?, seed: UInt64, - deEss: Bool + deEss: Bool, + temperature: Float ) async throws -> Data { logger.info("PocketTTS deterministic mode: seed=\(seed)") let session = try await makePocketSeededSession( - manager: manager, voice: voice, voiceData: voiceData, seed: seed) + manager: manager, voice: voice, voiceData: voiceData, seed: seed, + temperature: temperature) session.enqueue(text) session.finish() var allSamples: [Float] = [] @@ -668,17 +676,18 @@ public struct TTS { manager: PocketTtsManager, voice: String, voiceData: PocketTtsVoiceData?, - seed: UInt64 + seed: UInt64, + temperature: Float ) async throws -> PocketTtsSession { if let voiceData = voiceData { return try await manager.makeSession( voiceData: voiceData, - temperature: PocketTtsConstants.temperature, + temperature: temperature, seed: seed) } return try await manager.makeSession( voice: voice, - temperature: PocketTtsConstants.temperature, + temperature: temperature, seed: seed) } @@ -688,7 +697,8 @@ public struct TTS { voiceFilePath: String?, saveVoicePath: String?, language: PocketTtsLanguage, seed: UInt64? = nil, - placement: PocketTtsModelPlacement = .gpu + placement: PocketTtsModelPlacement = .gpu, + temperature: Float = PocketTtsConstants.temperature ) async { do { let tStart = Date() @@ -734,13 +744,14 @@ public struct TTS { voice: pocketVoice, voiceData: voiceData, seed: seed, - deEss: deEss) + deEss: deEss, + temperature: temperature) } else if let voiceData = voiceData { wav = try await manager.synthesize( - text: text, voiceData: voiceData, deEss: deEss) + text: text, voiceData: voiceData, temperature: temperature, deEss: deEss) } else { wav = try await manager.synthesize( - text: text, voice: pocketVoice, deEss: deEss) + text: text, voice: pocketVoice, temperature: temperature, deEss: deEss) } let tSynth1 = Date() @@ -1373,6 +1384,7 @@ public struct TTS { portuguese, portuguese_24l, spanish, spanish_24l Note: French is 24-layer only (no 6-layer pack upstream) --seed N Deterministic-mode seed (uses session API for fixed RNG) + --temperature T Generation temperature (default 0.7) --placement P Model placement: gpu (default), ane (rank-4 ANE models), ane-state (Trial 23 MLState multifunction pipeline; macOS 15+/iOS 18+, requires pocket_state.mlmodelc) diff --git a/Tests/FluidAudioTests/TTS/PocketTTS/PocketTtsDownloadFilterTests.swift b/Tests/FluidAudioTests/TTS/PocketTTS/PocketTtsDownloadFilterTests.swift index 327cb1864..63d6fe966 100644 --- a/Tests/FluidAudioTests/TTS/PocketTTS/PocketTtsDownloadFilterTests.swift +++ b/Tests/FluidAudioTests/TTS/PocketTTS/PocketTtsDownloadFilterTests.swift @@ -72,6 +72,16 @@ final class PocketTtsDownloadFilterTests: XCTestCase { XCTAssertFalse(skips("v2.1/english/manifest.json")) } + func testPackLocalCloningEncoderIsSkippedInPackDownloads() { + // The per-language voice-cloning encoder (`mimi_encoderv3.mlmodelc`, + // #793) is not part of any required model set — it is fetched lazily + // by `ensurePackMimiEncoder` on first clone, so pack downloads for + // synthesis-only users stay lean. + XCTAssertTrue(skips("v2.1/spanish_24l/mimi_encoderv3.mlmodelc")) + XCTAssertTrue(skips("v2.1/spanish_24l/mimi_encoderv3.mlmodelc/weights/weight.bin")) + XCTAssertTrue(skips("v2.1/spanish_24l/mimi_encoderv3.mlmodelc", placement: .ane)) + } + func testHistoricalExclusionsStillApply() { // .mlpackage sources are skipped even for a required model name. XCTAssertTrue(skips("v2.1/english/cond_prefill.mlpackage")) From 745a59b7ab7d99d2c9f1ed95ff178bf6288fcf20 Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Wed, 19 Aug 2026 15:01:01 -0400 Subject: [PATCH 2/2] fix(tts/pockettts): harden pack-encoder cache acceptance and cancellation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review fixes for ensurePackMimiEncoder: - Cache acceptance now goes through ModelCache.isCacheComplete instead of bare directory existence. An interrupted download deliberately leaves *.partial staging files for byte-range resume; accepting such a bundle failed MLModel(contentsOf:) later without ever taking the shared-encoder fallback (same class as #819). A bundle still incomplete after the resume attempt is cleared and re-downloaded once from scratch (stale staging file next to complete weights — resume can't repair it), mirroring ModelHub.loadWithRecovery semantics; verified locally by planting a weight.bin.partial: detect -> clear -> re-download -> clone proceeds pack-local, warm cache stays a no-network fast path. - Cancellation is no longer swallowed as an optional-download failure: the helper is throwing and rethrows via RetryPolicy.isCancellation, so a cancelled caller propagates instead of continuing into the fallback (or kicking off another download). Genuine availability/network failures still return nil and fall back. --- .../Assets/PocketTtsResourceDownloader.swift | 41 +++++++++++++++++-- .../Pipeline/PocketTtsModelStore.swift | 2 +- 2 files changed, 38 insertions(+), 5 deletions(-) diff --git a/Sources/FluidAudio/TTS/PocketTTS/Assets/PocketTtsResourceDownloader.swift b/Sources/FluidAudio/TTS/PocketTTS/Assets/PocketTtsResourceDownloader.swift index 24f9de8f7..24cf2aa5d 100644 --- a/Sources/FluidAudio/TTS/PocketTTS/Assets/PocketTtsResourceDownloader.swift +++ b/Sources/FluidAudio/TTS/PocketTTS/Assets/PocketTtsResourceDownloader.swift @@ -358,9 +358,19 @@ public enum PocketTtsResourceDownloader { /// returns `nil` when the encoder can't be fetched (offline, or not yet /// published for the pack); callers then fall back to the shared root /// encoder + reprojection path (#793). + /// + /// Cache acceptance goes through `ModelCache.isCacheComplete`, not bare + /// directory existence: an interrupted download deliberately leaves + /// `*.partial` staging files behind for byte-range resume, and accepting + /// such a bundle would fail `MLModel(contentsOf:)` later without ever + /// taking the shared-encoder fallback (same class as issue #819). + /// + /// Throws only on cancellation — a cancelled caller must not be routed + /// into the fallback (or trigger further downloads); genuine + /// availability/network failures return `nil`. public static func ensurePackMimiEncoder( language: PocketTtsLanguage, directory: URL? = nil - ) async -> URL? { + ) async throws -> URL? { do { let targetDir = try directory ?? cacheDirectory() let modelsDirectory = targetDir.appendingPathComponent( @@ -370,7 +380,7 @@ public enum PocketTtsResourceDownloader { "\(language.repoSubdirectory)/\(ModelNames.PocketTTS.mimiEncoderV3File)" let encoderPath = repoDir.appendingPathComponent(encoderSubpath) - if FileManager.default.fileExists(atPath: encoderPath.path) { + if ModelCache.isCacheComplete(at: repoDir, requiredFiles: [encoderSubpath]) { return encoderPath } @@ -385,14 +395,37 @@ public enum PocketTtsResourceDownloader { to: repoDir ) - guard FileManager.default.fileExists(atPath: encoderPath.path) else { + if ModelCache.isCacheComplete(at: repoDir, requiredFiles: [encoderSubpath]) { + return encoderPath + } + + // Still incomplete after a resume attempt: a genuinely interrupted + // transfer resumes its `.partial` and completes above, so what's + // left is a corrupt bundle the resume path can't repair (e.g. a + // stale staging file next to complete weights). Clear it and + // re-download once from scratch — same delete-and-retry semantics + // as `ModelHub.loadWithRecovery`. + logger.warning( + "Per-language Mimi encoder bundle for \(language.rawValue) is corrupt after " + + "resume; clearing and re-downloading once...") + try? FileManager.default.removeItem(at: encoderPath) + try await ModelHub.download( + .pocketTts, + subdirectory: encoderSubpath, + to: repoDir + ) + + guard ModelCache.isCacheComplete(at: repoDir, requiredFiles: [encoderSubpath]) else { logger.warning( - "Per-language Mimi encoder unavailable for \(language.rawValue); " + "Per-language Mimi encoder unavailable or incomplete for \(language.rawValue); " + "falling back to the shared encoder + reprojection (#793).") return nil } return encoderPath } catch { + if RetryPolicy.isCancellation(error) { + throw error + } logger.warning( "Failed to fetch per-language Mimi encoder for \(language.rawValue): " + "\(error.localizedDescription). Falling back to the shared encoder + " diff --git a/Sources/FluidAudio/TTS/PocketTTS/Pipeline/PocketTtsModelStore.swift b/Sources/FluidAudio/TTS/PocketTTS/Pipeline/PocketTtsModelStore.swift index 67775ac52..405adce40 100644 --- a/Sources/FluidAudio/TTS/PocketTTS/Pipeline/PocketTtsModelStore.swift +++ b/Sources/FluidAudio/TTS/PocketTTS/Pipeline/PocketTtsModelStore.swift @@ -397,7 +397,7 @@ public actor PocketTtsModelStore { config.computeUnits = .cpuAndGPU if language != .english, - let packURL = await PocketTtsResourceDownloader.ensurePackMimiEncoder( + let packURL = try await PocketTtsResourceDownloader.ensurePackMimiEncoder( language: language, directory: directory) { logger.info("Loading pack-local Mimi encoder (\(self.language.rawValue))...")