Add Voice Cloning (CoreML)

Examples/TTS/TTSKitExample/TTSKitExample/ViewModel.swift

@@ -1198,8 +1198,8 @@ final class ViewModel: @unchecked Sendable {
 extension TTSModelVariant {
     var sizeEstimate: String {
         switch self {
-            case .qwen3TTS_0_6b: return "~1 GB"
-            case .qwen3TTS_1_7b: return "~2.2 GB"
+            case .qwen3TTS_0_6b, .qwen3TTS_0_6b_base: return "~1 GB"
+            case .qwen3TTS_1_7b, .qwen3TTS_1_7b_base: return "~2.2 GB"
         }
     }
 }

Makefile

@@ -9,7 +9,7 @@ PYTHON_COMMAND := python3
 MODEL_REPO := argmaxinc/whisperkit-coreml
 MODEL_REPO_DIR := ./Models/whisperkit-coreml
 TTS_MODEL_REPO := argmaxinc/ttskit-coreml
-TTS_MODEL_REPO_DIR := ./Models/ttskit-coreml
+TTS_MODEL_REPO_DIR := ./Models/$(notdir $(TTS_MODEL_REPO))
 SPEAKERKIT_MODEL_REPO := argmaxinc/speakerkit-coreml
 SPEAKERKIT_MODEL_REPO_DIR := ./Models/speakerkit-coreml
 BASE_COMPILED_DIR := ./Models
@@ -144,15 +144,23 @@ download-tts-models: setup-tts-model-repo
 # Download a specific TTS model size
 # Usage: make download-tts-model MODEL=0.6b
 #        make download-tts-model MODEL=1.7b
+#        make download-tts-model MODEL=0.6b-base
+# Base variants map to the 12hz-<MODEL> version dir and include the voice-clone
+# encoders (speaker_encoder, speech_encoder, speech_encoder_rvq) via the
+# component wildcard. Override the source repo with TTS_MODEL_REPO=<org>/<repo>.
 download-tts-model: setup-tts-model-repo
 	@if [ -z "$(MODEL)" ]; then \
 		echo "Error: MODEL not set. Usage: make download-tts-model MODEL=0.6b"; \
-		echo "Available models: 0.6b, 1.7b"; \
+		echo "Available models: 0.6b, 1.7b, 0.6b-base"; \
 		exit 1; \
 	fi
 	@echo "Downloading TTS model $(MODEL)..."
-	@cd $(TTS_MODEL_REPO_DIR) && \
-	git lfs pull --include="qwen3_tts/*/12hz-$(MODEL)-customvoice/**"
+	@case "$(MODEL)" in \
+		*-base) VERSION_DIR="12hz-$(MODEL)";; \
+		*) VERSION_DIR="12hz-$(MODEL)-customvoice";; \
+	esac; \
+	cd $(TTS_MODEL_REPO_DIR) && \
+	git lfs pull --include="qwen3_tts/*/$$VERSION_DIR/**"
 
 build:
 	@echo "Building argmax-oss-swift..."

Package.swift

@@ -110,6 +110,7 @@ let package = Package(
             dependencies: [
                 "TTSKit"
             ],
+            resources: [.copy("Resources")],
             swiftSettings: swiftSettings()
         ),
         .testTarget(

Package@swift-6.2.swift

@@ -114,6 +114,7 @@ let package = Package(
             dependencies: [
                 "TTSKit"
             ],
+            resources: [.copy("Resources")],
             swiftSettings: swiftSettings(libraryEvolution: false)
         ),
         .testTarget(

Sources/ArgmaxCLI/TTSCLI.swift

@@ -76,10 +76,21 @@ struct TTSCLI: AsyncParsableCommand {
     @Option(name: .long, help: "Random seed for reproducible output")
     var seed: UInt64?
 
+    // MARK: - Voice clone options
+
+    @Option(name: .long, help: "Reference audio clip to clone the voice from (any readable audio format; enables voice cloning)")
+    var refAudio: String?
+
+    @Option(name: .long, help: "Transcript of the reference clip (required with --ref-audio unless --x-vector-only)")
+    var refText: String?
+
+    @Flag(name: .long, help: "Clone with the speaker x-vector only, skipping reference RVQ encoding (lower fidelity, no --ref-text needed)")
+    var xVectorOnly: Bool = false
+
     // MARK: - Model selection
 
-    @Option(name: .long, help: "Model preset (0.6b, 1.7b). Auto-configures version dir and variant defaults.")
-    var model: TTSModelVariant = .qwen3TTS_0_6b
+    @Option(name: .long, help: "Model preset (0.6b, 0.6b-base, 1.7b, 1.7b-base). Auto-configures version dir and variant defaults; the -base presets carry the voice-clone assets. Defaults to 0.6b, or 0.6b-base when --ref-audio is set.")
+    var model: TTSModelVariant?
 
     // MARK: - Advanced options (auto-configured by preset, can be overridden)
 
@@ -110,8 +121,26 @@ struct TTSCLI: AsyncParsableCommand {
     @Option(name: .long, help: "SpeechDecoder variant (overrides --model preset)")
     var speechDecoderVariant: String?
 
-    @Option(name: .long, help: "SpeechDecoder mode: latencyOptimized (default, lowest time-to-first-audio, 1 frame/call) or throughputOptimized (higher throughput, ~4x larger pre-buffer, 4 frames/call)")
-    var speechDecoderMode: Qwen3SpeechDecoderMode = .latencyOptimized
+    @Option(name: .long, help: "CodeEmbedder variant (overrides --model preset)")
+    var codeEmbedderVariant: String?
+
+    @Option(name: .long, help: "MultiCodeEmbedder variant (overrides --model preset)")
+    var multiCodeEmbedderVariant: String?
+
+    @Option(name: .long, help: "TextProjector variant (overrides --model preset)")
+    var textProjectorVariant: String?
+
+    @Option(name: .long, help: "SpeakerEncoder variant for voice cloning (e.g. W16A16-15s for longer references)")
+    var speakerEncoderVariant: String?
+
+    @Option(name: .long, help: "SpeechEncoder variant for voice cloning (e.g. W16A16-15s for longer references)")
+    var speechEncoderVariant: String?
+
+    @Option(name: .long, help: "SpeechEncoderRVQ variant for voice cloning (must match --speech-encoder-variant window)")
+    var speechEncoderRVQVariant: String?
+
+    @Option(name: .long, help: "SpeechDecoder mode: latencyOptimized (lowest time-to-first-audio, 1 frame/call), throughputOptimized (higher throughput, ~4x larger pre-buffer, 4 frames/call), or singleFunction (single-function assets, e.g. the base-family speech decoders). Defaults to latencyOptimized, or singleFunction for -base model presets.")
+    var speechDecoderMode: Qwen3SpeechDecoderMode?
 
     // MARK: - Compute unit options
 
@@ -149,6 +178,22 @@ struct TTSCLI: AsyncParsableCommand {
             throw ValidationError("Input text is empty")
         }
 
+        // Validate voice-clone flag combinations.
+        if refAudio == nil {
+            if refText != nil {
+                throw ValidationError("--ref-text requires --ref-audio")
+            }
+            if xVectorOnly {
+                throw ValidationError("--x-vector-only requires --ref-audio")
+            }
+        } else if !xVectorOnly, refText == nil {
+            throw ValidationError("--ref-text is required with --ref-audio (unless --x-vector-only is set)")
+        }
+
+        // Voice cloning needs the base-family checkpoints; default to 0.6b-base
+        // when --ref-audio is set and no explicit --model was given.
+        let model = self.model ?? (refAudio != nil ? .qwen3TTS_0_6b_base : .qwen3TTS_0_6b)
+
         // Resolve local models path if provided
         let resolvedModelFolder: URL? = modelsPath.map {
             URL(fileURLWithPath: FileManager.resolveAbsolutePath($0))
@@ -169,8 +214,14 @@ struct TTSCLI: AsyncParsableCommand {
             versionDir: versionDir,
             codeDecoderVariant: codeDecoderVariant,
             multiCodeDecoderVariant: multiCodeDecoderVariant,
+            codeEmbedderVariant: codeEmbedderVariant,
+            multiCodeEmbedderVariant: multiCodeEmbedderVariant,
+            textProjectorVariant: textProjectorVariant,
             speechDecoderVariant: speechDecoderVariant,
-            speechDecoderMode: speechDecoderMode,
+            speakerEncoderVariant: speakerEncoderVariant,
+            speechEncoderVariant: speechEncoderVariant,
+            speechEncoderRVQVariant: speechEncoderRVQVariant,
+            speechDecoderMode: speechDecoderMode ?? (model.isBaseVariant ? .singleFunction : .latencyOptimized),
             computeOptions: ComputeOptions(
                 embedderComputeUnits: embedderComputeUnits.asMLComputeUnits,
                 codeDecoderComputeUnits: codeDecoderComputeUnits.asMLComputeUnits,
@@ -180,6 +231,13 @@ struct TTSCLI: AsyncParsableCommand {
             verbose: verbose
         )
 
+        // Voice cloning needs the three encoder assets, which sit outside the
+        // default component download patterns (they are loaded lazily by
+        // `loadVoiceCloneModels()`); include them in the model download.
+        if refAudio != nil {
+            config.downloadAdditionalPatterns += config.voiceCloneDownloadPatterns
+        }
+
         // Default: --play uses sequential (1), file output uses unlimited (0).
         let effectiveWorkerCount = concurrentWorkerCount ?? (play ? 1 : 0)
 
@@ -188,7 +246,7 @@ struct TTSCLI: AsyncParsableCommand {
 
         // Warn if instruction is used with a model that doesn't support it
         var effectiveInstruction = instruction
-        if let instruction = effectiveInstruction, !instruction.isEmpty, model == .qwen3TTS_0_6b {
+        if let instruction = effectiveInstruction, !instruction.isEmpty, !model.supportsVoiceDirection {
             print("Warning: --instruction is only supported by the 1.7B model variant. Ignoring instruction for \(model.rawValue).")
             effectiveInstruction = nil
         }
@@ -202,6 +260,13 @@ struct TTSCLI: AsyncParsableCommand {
             print("  Speaker: \(speaker.rawValue)")
             print("  Language: \(language.rawValue)")
             print("  Model: \(model.rawValue)")
+            if let refAudio {
+                print("  Reference audio: \(refAudio)")
+                if let refText {
+                    print("  Reference text: \"\(refText.prefix(80))\(refText.count > 80 ? "..." : "")\"")
+                }
+                print("  Clone mode: \(xVectorOnly ? "x-vector only" : "ICL")")
+            }
             if let inst = effectiveInstruction {
                 print("  Instruction: \"\(inst)\"")
             }
@@ -231,6 +296,19 @@ struct TTSCLI: AsyncParsableCommand {
         config.seed = effectiveSeed
         let tts = try await TTSKit(config)
 
+        // Encode the reference clip once; the resulting prompt is reused for
+        // every text chunk of this generate call.
+        var voiceClonePrompt: VoiceClonePrompt?
+        if let refAudio {
+            let refURL = URL(fileURLWithPath: FileManager.resolveAbsolutePath(refAudio))
+            try await tts.loadVoiceCloneModels()
+            voiceClonePrompt = try await tts.cloneVoice(
+                referenceAudio: refURL,
+                referenceText: refText,
+                xVectorOnly: xVectorOnly
+            )
+        }
+
         let options = GenerationOptions(
             temperature: temperature,
             topK: topK,
@@ -239,7 +317,8 @@ struct TTSCLI: AsyncParsableCommand {
             concurrentWorkerCount: effectiveWorkerCount,
             targetChunkSize: targetChunkSize,
             minChunkSize: minChunkSize,
-            instruction: effectiveInstruction
+            instruction: effectiveInstruction,
+            voiceClone: voiceClonePrompt
         )
 
         let result: SpeechResult

Sources/TTSKit/Models.swift

@@ -251,6 +251,8 @@ public struct GenerationOptions: Codable, Sendable {
     /// (e.g., `"Very happy"`). Prepended as a text-only user prompt before the main
     /// TTS segment. For Qwen3, this is only supported by the 1.7B model variant.
     public var instruction: String?
+    /// Voice-clone prompt derived from a reference clip; nil for custom-voice TTS.
+    public var voiceClone: VoiceClonePrompt?
 
     /// Force the legacy `[FloatType]` inference path even on macOS 15+ / iOS 18+.
     /// When `false` (default), the MLTensor path is taken on supported OS versions.
@@ -268,6 +270,7 @@ public struct GenerationOptions: Codable, Sendable {
         targetChunkSize: Int? = nil,
         minChunkSize: Int? = nil,
         instruction: String? = nil,
+        voiceClone: VoiceClonePrompt? = nil,
         forceLegacyEmbedPath: Bool = false
     ) {
         self.temperature = temperature
@@ -279,6 +282,7 @@ public struct GenerationOptions: Codable, Sendable {
         self.targetChunkSize = targetChunkSize
         self.minChunkSize = minChunkSize
         self.instruction = instruction
+        self.voiceClone = voiceClone
         self.forceLegacyEmbedPath = forceLegacyEmbedPath
     }
 }