add auto download for the missing mlmodelc, phoene mappers & vocabs

Author: Alex-Wengg

Sources/FluidAudio/TextToSpeech/Kokoro/Assets/TtsResourceDownloader.swift

@@ -4,21 +4,36 @@ import Foundation
 public enum TtsResourceDownloader {
 
     private static let logger = AppLogger(category: "TtsResourceDownloader")
+    private static let zhRepoRemotePath = "alexwengg/tts-zh"
 
     /// Download a voice embedding JSON file from HuggingFace
     public static func downloadVoiceEmbedding(voice: String) async throws -> Data {
-        let url = try ModelRegistry.resolveModel(Repo.kokoro.remotePath, "voices/\(voice).json")
+        // Primary repo (FluidInference/kokoro-82m-coreml)
+        let primary = try ModelRegistry.resolveModel(Repo.kokoro.remotePath, "voices/\(voice).json")
+        // Fallback repo for zh voices
+        let fallback = try ModelRegistry.resolveModel(zhRepoRemotePath, "voices/\(voice).json")
 
-        do {
-            let data = try await AssetDownloader.fetchData(
+        func attempt(_ url: URL) async throws -> Data {
+            try await AssetDownloader.fetchData(
                 from: url,
                 description: "\(voice) voice embedding JSON",
                 logger: logger
             )
-            logger.info("Downloaded voice embedding JSON for \(voice)")
+        }
+
+        do {
+            let data = try await attempt(primary)
+            logger.info("Downloaded voice embedding JSON for \(voice) from primary repo")
             return data
         } catch {
-            throw TTSError.modelNotFound("Voice embedding JSON unavailable for \(voice): \(error.localizedDescription)")
+            logger.warning("Primary voice embedding missing for \(voice); trying zh repo")
+            do {
+                let data = try await attempt(fallback)
+                logger.info("Downloaded voice embedding JSON for \(voice) from zh repo")
+                return data
+            } catch {
+                throw TTSError.modelNotFound("Voice embedding JSON unavailable for \(voice)")
+            }
         }
     }
 
@@ -74,3 +89,71 @@ public enum TtsResourceDownloader {
     }
 
 }
+
+// MARK: - zh assets (vocab + char lexicon)
+
+extension TtsResourceDownloader {
+    /// Ensure `zh_vocab_index.json` exists in Kokoro cache. Falls back to downloading `zh.json` and writing it
+    /// as `zh_vocab_index.json` if needed.
+    @discardableResult
+    public static func ensureZhVocabularyInCache() async throws -> URL {
+        let cacheDir = try TtsModels.cacheDirectoryURL()
+        let kokoroDir = cacheDir.appendingPathComponent("Models/kokoro")
+        try FileManager.default.createDirectory(at: kokoroDir, withIntermediateDirectories: true)
+
+        let zhVocabURL = kokoroDir.appendingPathComponent("zh_vocab_index.json")
+        if FileManager.default.fileExists(atPath: zhVocabURL.path) {
+            return zhVocabURL
+        }
+
+        let primary = try ModelRegistry.resolveModel(zhRepoRemotePath, "zh_vocab_index.json")
+        let alt = try ModelRegistry.resolveModel(zhRepoRemotePath, "zh.json")
+
+        do {
+            let descriptor = AssetDownloader.Descriptor(
+                description: "zh_vocab_index.json",
+                remoteURL: primary,
+                destinationURL: zhVocabURL
+            )
+            return try await AssetDownloader.ensure(descriptor, logger: logger)
+        } catch {
+            logger.warning("zh_vocab_index.json not found in zh repo; trying zh.json")
+            let data = try await AssetDownloader.fetchData(
+                from: alt,
+                description: "zh.json",
+                logger: logger
+            )
+            try data.write(to: zhVocabURL, options: [.atomic])
+            logger.info("Cached zh vocabulary as zh_vocab_index.json")
+            return zhVocabURL
+        }
+    }
+
+    /// Ensure `zh_char_phonemes.json` exists in Kokoro cache, downloading from zh repo if missing.
+    @discardableResult
+    public static func ensureZhCharPhonemesInCache() async throws -> URL {
+        let cacheDir = try TtsModels.cacheDirectoryURL()
+        let kokoroDir = cacheDir.appendingPathComponent("Models/kokoro")
+        try FileManager.default.createDirectory(at: kokoroDir, withIntermediateDirectories: true)
+
+        let localURL = kokoroDir.appendingPathComponent("zh_char_phonemes.json")
+        if FileManager.default.fileExists(atPath: localURL.path) {
+            return localURL
+        }
+
+        let remoteURL = try ModelRegistry.resolveModel(zhRepoRemotePath, "zh_char_phonemes.json")
+        let descriptor = AssetDownloader.Descriptor(
+            description: "zh_char_phonemes.json",
+            remoteURL: remoteURL,
+            destinationURL: localURL
+        )
+        return try await AssetDownloader.ensure(descriptor, logger: logger)
+    }
+
+    /// Convenience: ensure both zh vocab and char phoneme lexicon exist in cache.
+    public static func ensureZhAssetsInCache() async throws -> (vocabURL: URL, lexURL: URL) {
+        async let v = ensureZhVocabularyInCache()
+        async let l = ensureZhCharPhonemesInCache()
+        return try await (vocabURL: v, lexURL: l)
+    }
+}

Sources/FluidAudioCLI/Commands/TTSZhCommand.swift

@@ -92,6 +92,14 @@ enum TTSZh {
                 }
             }
 
+            // Ensure zh assets (vocab + char lexicon) in cache as a fallback, and prefer zh vocab override.
+            do {
+                let ensured = try await TtsResourceDownloader.ensureZhAssetsInCache()
+                await KokoroVocabulary.shared.setOverrideURL(ensured.vocabURL)
+            } catch {
+                logger.warning("Failed to ensure zh assets; continuing: \(error.localizedDescription)")
+            }
+
             // Provide extra voices dirs: explicit + sibling voices
             var extraDirs: [URL] = []
             if let dir = voicesDir, !dir.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty {