fix(tts/pockettts): per-language mimi encoders for non-English voice cloning (#793)

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/<lang>/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"

Sources/FluidAudio/TTS/PocketTTS/Assets/PocketTtsResourceDownloader.swift

@@ -349,6 +349,91 @@ public enum PocketTtsResourceDownloader {
         return encoderPath
     }
 
+    /// Ensure the pack-local per-language voice-cloning encoder
+    /// (`v2.1/<lang>/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).
+    ///
+    /// 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 throws -> 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 ModelCache.isCacheComplete(at: repoDir, requiredFiles: [encoderSubpath]) {
+                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
+            )
+
+            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 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 + "
+                    + "reprojection (#793).")
+            return nil
+        }
+    }
+
     /// Ensure voice conditioning data for the given language is available,
     /// downloading from HuggingFace if missing.
     ///

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 = try 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/<lang>/` 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/<lang>/`
+        // 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) "

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)

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"))