Add swift test CLI

.gitignore

@@ -3,3 +3,4 @@
 __pycache__
 test_data
 **.wav
+swift-cli/.build

swift-cli/Package.resolved

@@ -0,0 +1,104 @@
+{
+  "pins" : [
+    {
+      "identity" : "eventsource",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/mattt/EventSource.git",
+      "state" : {
+        "revision" : "a3a85a85214caf642abaa96ae664e4c772a59f6e",
+        "version" : "1.4.1"
+      }
+    },
+    {
+      "identity" : "swift-asn1",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-asn1.git",
+      "state" : {
+        "revision" : "9f542610331815e29cc3821d3b6f488db8715517",
+        "version" : "1.6.0"
+      }
+    },
+    {
+      "identity" : "swift-atomics",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-atomics.git",
+      "state" : {
+        "revision" : "b601256eab081c0f92f059e12818ac1d4f178ff7",
+        "version" : "1.3.0"
+      }
+    },
+    {
+      "identity" : "swift-collections",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-collections.git",
+      "state" : {
+        "revision" : "6675bc0ff86e61436e615df6fc5174e043e57924",
+        "version" : "1.4.1"
+      }
+    },
+    {
+      "identity" : "swift-crypto",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-crypto.git",
+      "state" : {
+        "revision" : "bb4ba815dab96d4edc1e0b86d7b9acf9ff973a84",
+        "version" : "4.3.1"
+      }
+    },
+    {
+      "identity" : "swift-huggingface",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/huggingface/swift-huggingface.git",
+      "state" : {
+        "revision" : "b721959445b617d0bf03910b2b4aced345fd93bf",
+        "version" : "0.9.0"
+      }
+    },
+    {
+      "identity" : "swift-jinja",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/huggingface/swift-jinja.git",
+      "state" : {
+        "revision" : "0aeefadec459ce8e11a333769950fb86183aca43",
+        "version" : "2.3.5"
+      }
+    },
+    {
+      "identity" : "swift-nio",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-nio.git",
+      "state" : {
+        "revision" : "558f24a4647193b5a0e2104031b71c55d31ff83a",
+        "version" : "2.97.1"
+      }
+    },
+    {
+      "identity" : "swift-system",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-system.git",
+      "state" : {
+        "revision" : "7c6ad0fc39d0763e0b699210e4124afd5041c5df",
+        "version" : "1.6.4"
+      }
+    },
+    {
+      "identity" : "swift-transformers",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/huggingface/swift-transformers",
+      "state" : {
+        "revision" : "b38443e44d93eca770f2eb68e2a4d0fa100f9aa2",
+        "version" : "1.3.0"
+      }
+    },
+    {
+      "identity" : "yyjson",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/ibireme/yyjson.git",
+      "state" : {
+        "revision" : "8b4a38dc994a110abaec8a400615567bd996105f",
+        "version" : "0.12.0"
+      }
+    }
+  ],
+  "version" : 2
+}

swift-cli/Package.swift

@@ -0,0 +1,19 @@
+// swift-tools-version: 5.9
+import PackageDescription
+
+let package = Package(
+    name: "plapre-cli",
+    platforms: [.macOS("15.0")],
+    dependencies: [
+        .package(url: "https://github.com/huggingface/swift-transformers", from: "1.3.0"),
+    ],
+    targets: [
+        .executableTarget(
+            name: "plapre-cli",
+            dependencies: [
+                .product(name: "Tokenizers", package: "swift-transformers"),
+            ],
+            path: "Sources"
+        ),
+    ]
+)

swift-cli/Sources/main.swift

@@ -0,0 +1,714 @@
+import Foundation
+import CoreML
+import Accelerate
+import Tokenizers
+
+// MARK: - Constants
+
+let sampleRate: Int = 24000
+let prefillSeqLen = 512
+let maxContext = 2048
+let headDim = 64
+let numKvHeads = 3
+let speakerDim = 128
+let audioTokenOffset = 8002
+let audioMarkerToken: Int32 = 8001
+let textMarkerToken: Int32 = 8000
+let eosToken: Int32 = 0  // <eos> is token 0 in plapre tokenizer
+let vocabSize = 20802
+
+// HiFT source generation parameters
+let hiftNfft = 16
+let hiftHopLen = 4
+let hiftSamplingRate: Float = 24000.0
+let hiftHarmonicNum = 8
+let hiftSineAmp: Float = 0.1
+let hiftNoiseStd: Float = 0.003
+let hiftUpsampleScale = 480
+let hiftWindow: [Float] = [0.0, 0.03806023, 0.14644662, 0.30865827, 0.5, 0.6913417, 0.85355341, 0.96193975, 1.0, 0.96193975, 0.85355341, 0.6913417, 0.5, 0.30865827, 0.14644662, 0.03806023]
+// l_linear: 9 harmonics → 1, then tanh
+let sourceLinearWeight: [Float] = [-0.27458203, -0.27744064, 0.07214482, 0.12596518, 0.02788151, 0.00307915, 0.01020926, -0.01141518, -0.01324173]
+let sourceLinearBias: Float = 7.7338242e-05
+
+// MARK: - Model paths
+
+let repoRoot = URL(fileURLWithPath: #filePath)
+    .deletingLastPathComponent()  // Sources
+    .deletingLastPathComponent()  // swift-cli
+    .deletingLastPathComponent()  // repo root
+
+func modelURL(_ name: String) -> URL {
+    repoRoot.appendingPathComponent("\(name).mlpackage")
+}
+
+// MARK: - RoPE tables
+
+func loadRopeTable(_ name: String) -> [Float] {
+    let url = repoRoot.appendingPathComponent(name)
+    // .npy format: 128-byte header + raw float16 data
+    let data = try! Data(contentsOf: url)
+    // Find header end (newline after header)
+    var headerEnd = 0
+    for i in 0..<data.count {
+        if data[i] == 0x0A {
+            // Check if this could be the end of a npy header
+            // npy header ends with \n, and the header size is padded to multiple of 64
+            if i > 5 {
+                headerEnd = i + 1
+                // Verify the remaining data makes sense
+                let remaining = data.count - headerEnd
+                if remaining % 2 == 0 { // float16 = 2 bytes
+                    break
+                }
+            }
+        }
+    }
+    let rawData = data.subdata(in: headerEnd..<data.count)
+    // Convert float16 to float32
+    let count = rawData.count / 2
+    var result = [Float](repeating: 0, count: count)
+    rawData.withUnsafeBytes { ptr in
+        let f16 = ptr.bindMemory(to: Float16.self)
+        for i in 0..<count {
+            result[i] = Float(f16[i])
+        }
+    }
+    return result
+}
+
+// MARK: - Speaker embeddings
+
+func loadSpeaker(_ name: String) -> [Float] {
+    let url = repoRoot.appendingPathComponent("speakers.json")
+    let data = try! Data(contentsOf: url)
+    let json = try! JSONSerialization.jsonObject(with: data) as! [String: [Double]]
+    guard let emb = json[name] else {
+        fatalError("Speaker '\(name)' not found. Available: \(json.keys.sorted())")
+    }
+    return emb.map { Float($0) }
+}
+
+
+// MARK: - Tokenizer
+
+// MARK: - CoreML helpers
+
+func compileModel(at url: URL) throws -> MLModel {
+    print("  Compiling \(url.lastPathComponent)...")
+    let compiled = try MLModel.compileModel(at: url)
+    let config = MLModelConfiguration()
+    config.computeUnits = .cpuOnly
+    return try MLModel(contentsOf: compiled, configuration: config)
+}
+
+func mlArray(_ values: [Float], shape: [Int]) -> MLMultiArray {
+    let arr = try! MLMultiArray(shape: shape.map { NSNumber(value: $0) }, dataType: .float16)
+    let count = values.count
+    arr.withUnsafeMutableBufferPointer(ofType: Float16.self) { ptr, _ in
+        for i in 0..<count {
+            ptr[i] = Float16(values[i])
+        }
+    }
+    return arr
+}
+
+func mlArrayFloat32(_ values: [Float], shape: [Int]) -> MLMultiArray {
+    let arr = try! MLMultiArray(shape: shape.map { NSNumber(value: $0) }, dataType: .float32)
+    arr.withUnsafeMutableBufferPointer(ofType: Float.self) { dst, _ in
+        for i in 0..<values.count {
+            dst[i] = values[i]
+        }
+    }
+    return arr
+}
+
+func mlArrayInt32(_ values: [Int32], shape: [Int]) -> MLMultiArray {
+    let arr = try! MLMultiArray(shape: shape.map { NSNumber(value: $0) }, dataType: .int32)
+    arr.withUnsafeMutableBufferPointer(ofType: Int32.self) { dst, _ in
+        for i in 0..<values.count {
+            dst[i] = values[i]
+        }
+    }
+    return arr
+}
+
+func readFloat16Array(_ arr: MLMultiArray) -> [Float] {
+    let count = arr.count
+    var result = [Float](repeating: 0, count: count)
+    arr.withUnsafeBufferPointer(ofType: Float16.self) { ptr in
+        for i in 0..<count {
+            result[i] = Float(ptr[i])
+        }
+    }
+    return result
+}
+
+func readFloat32Array(_ arr: MLMultiArray) -> [Float] {
+    let count = arr.count
+    var result = [Float](repeating: 0, count: count)
+    arr.withUnsafeBufferPointer(ofType: Float.self) { ptr in
+        for i in 0..<count {
+            result[i] = ptr[i]
+        }
+    }
+    return result
+}
+
+// MARK: - Source signal generation (replaces HiFT's m_source in Swift)
+
+func generateSourceSTFT(f0: [Float], melLength: Int) -> [Float] {
+    // f0 shape: (melLength,) — one f0 value per mel frame
+
+    // 1. Upsample f0 by hiftUpsampleScale (nearest neighbor)
+    let audioLength = melLength * hiftUpsampleScale
+    var f0Up = [Float](repeating: 0, count: audioLength)
+    for i in 0..<audioLength {
+        f0Up[i] = f0[min(i / hiftUpsampleScale, melLength - 1)]
+    }
+
+    // 2. Generate harmonics: f0 * [1, 2, ..., harmonic_num+1]
+    let numHarmonics = hiftHarmonicNum + 1  // 9
+    var sineWaves = [[Float]](repeating: [Float](repeating: 0, count: audioLength), count: numHarmonics)
+
+    for h in 0..<numHarmonics {
+        let harmonicMul = Float(h + 1)
+        // Cumulative phase: phase[t] = sum(f0[0..t] * harmonic / sr) * 2pi
+        var phase: Float = 0
+        for t in 0..<audioLength {
+            let f = f0Up[t] * harmonicMul
+            phase += f / hiftSamplingRate
+            // Keep phase in [0, 1) to avoid precision loss
+            phase = phase - Float(Int(phase))
+            sineWaves[h][t] = sin(phase * 2 * .pi) * hiftSineAmp
+        }
+    }
+
+    // 3. UV detection: voiced (f0 > 0) vs unvoiced
+    var uv = [Float](repeating: 0, count: audioLength)
+    for t in 0..<audioLength {
+        uv[t] = f0Up[t] > 0 ? 1.0 : 0.0
+    }
+
+    // 4. Apply UV masking + noise
+    for h in 0..<numHarmonics {
+        for t in 0..<audioLength {
+            let noise = Float.random(in: -1...1) * (uv[t] * hiftNoiseStd + (1 - uv[t]) * hiftSineAmp / 3)
+            sineWaves[h][t] = sineWaves[h][t] * uv[t] + noise
+        }
+    }
+
+    // 5. Linear combination: 9 harmonics → 1 via l_linear + tanh
+    var source = [Float](repeating: 0, count: audioLength)
+    for t in 0..<audioLength {
+        var val: Float = sourceLinearBias
+        for h in 0..<numHarmonics {
+            val += sineWaves[h][t] * sourceLinearWeight[h]
+        }
+        source[t] = tanh(val)
+    }
+
+    // 6. STFT of source signal
+    // n_fft=16, hop=4, hann window
+    let nfftHalf = hiftNfft / 2 + 1  // 9
+    let numFrames = audioLength / hiftHopLen + 1
+    // Output: (18, numFrames) — 9 real + 9 imag channels
+    var stft = [Float](repeating: 0, count: 18 * numFrames)
+
+    for frame in 0..<numFrames {
+        let center = frame * hiftHopLen
+        // Windowed segment
+        var segment = [Float](repeating: 0, count: hiftNfft)
+        for k in 0..<hiftNfft {
+            let idx = center - hiftNfft / 2 + k
+            if idx >= 0 && idx < audioLength {
+                segment[k] = source[idx] * hiftWindow[k]
+            }
+        }
+        // DFT for each frequency bin
+        for f in 0..<nfftHalf {
+            var real: Float = 0
+            var imag: Float = 0
+            for k in 0..<hiftNfft {
+                let angle = -2.0 * Float.pi * Float(f) * Float(k) / Float(hiftNfft)
+                real += segment[k] * cos(angle)
+                imag += segment[k] * sin(angle)
+            }
+            stft[f * numFrames + frame] = real          // real part
+            stft[(nfftHalf + f) * numFrames + frame] = imag  // imag part
+        }
+    }
+
+    return stft
+}
+
+// MARK: - iSTFT (magnitude + phase → waveform)
+
+func istft(magnitude: [Float], phase: [Float], numFrames: Int) -> [Float] {
+    // Matches torch.istft(spec, n_fft=16, hop_length=4, win_length=16, window=hann, center=True)
+    let nfftHalf = hiftNfft / 2 + 1  // 9
+    // center=True means the STFT was padded by n_fft//2 on each side
+    // Total overlap-add length includes this padding
+    let padded_length = (numFrames - 1) * hiftHopLen + hiftNfft
+    var output = [Float](repeating: 0, count: padded_length)
+    var windowSum = [Float](repeating: 0, count: padded_length)
+
+    for frame in 0..<numFrames {
+        // Build full complex spectrum from one-sided
+        var real = [Float](repeating: 0, count: hiftNfft)
+        var imag = [Float](repeating: 0, count: hiftNfft)
+
+        for f in 0..<nfftHalf {
+            let mag = magnitude[f * numFrames + frame]
+            let ph = phase[f * numFrames + frame]
+            real[f] = mag * cos(ph)
+            imag[f] = mag * sin(ph)
+        }
+        // Mirror for negative frequencies (Hermitian symmetry)
+        for f in 1..<(hiftNfft / 2) {
+            real[hiftNfft - f] = real[f]
+            imag[hiftNfft - f] = -imag[f]
+        }
+
+        // IDFT
+        var segment = [Float](repeating: 0, count: hiftNfft)
+        for k in 0..<hiftNfft {
+            var val: Float = 0
+            for fi in 0..<hiftNfft {
+                let angle = 2.0 * Float.pi * Float(fi) * Float(k) / Float(hiftNfft)
+                val += real[fi] * cos(angle) - imag[fi] * sin(angle)
+            }
+            segment[k] = val / Float(hiftNfft)
+        }
+
+        // Overlap-add with window
+        let start = frame * hiftHopLen
+        for k in 0..<hiftNfft {
+            let idx = start + k
+            if idx < padded_length {
+                output[idx] += segment[k] * hiftWindow[k]
+                windowSum[idx] += hiftWindow[k] * hiftWindow[k]
+            }
+        }
+    }
+
+    // Normalize by window sum
+    for i in 0..<padded_length {
+        if windowSum[i] > 1e-8 {
+            output[i] /= windowSum[i]
+        }
+    }
+
+    // Trim center padding: remove n_fft//2 from start, and from end to match expected length
+    let pad = hiftNfft / 2  // 8
+    let expectedLength = (numFrames - 1) * hiftHopLen  // what torch.istft returns
+    let trimStart = pad
+    let trimEnd = min(trimStart + expectedLength, padded_length)
+    var trimmed = Array(output[trimStart..<trimEnd])
+
+    // Clamp
+    for i in 0..<trimmed.count {
+        trimmed[i] = max(-0.99, min(0.99, trimmed[i]))
+    }
+
+    return trimmed
+}
+
+// MARK: - WAV writer
+
+func writeWAV(_ samples: [Float], to url: URL, sampleRate: Int = 24000) {
+    let numSamples = samples.count
+    let dataSize = numSamples * 2  // 16-bit PCM
+    var data = Data()
+
+    // RIFF header
+    data.append(contentsOf: "RIFF".utf8)
+    var chunkSize = UInt32(36 + dataSize).littleEndian
+    data.append(Data(bytes: &chunkSize, count: 4))
+    data.append(contentsOf: "WAVE".utf8)
+
+    // fmt chunk
+    data.append(contentsOf: "fmt ".utf8)
+    var fmtSize = UInt32(16).littleEndian; data.append(Data(bytes: &fmtSize, count: 4))
+    var audioFormat = UInt16(1).littleEndian; data.append(Data(bytes: &audioFormat, count: 2))
+    var channels = UInt16(1).littleEndian; data.append(Data(bytes: &channels, count: 2))
+    var sr = UInt32(sampleRate).littleEndian; data.append(Data(bytes: &sr, count: 4))
+    var byteRate = UInt32(sampleRate * 2).littleEndian; data.append(Data(bytes: &byteRate, count: 4))
+    var blockAlign = UInt16(2).littleEndian; data.append(Data(bytes: &blockAlign, count: 2))
+    var bitsPerSample = UInt16(16).littleEndian; data.append(Data(bytes: &bitsPerSample, count: 2))
+
+    // data chunk
+    data.append(contentsOf: "data".utf8)
+    var dataChunkSize = UInt32(dataSize).littleEndian; data.append(Data(bytes: &dataChunkSize, count: 4))
+    for s in samples {
+        let clamped = max(-1.0, min(1.0, s))
+        var pcm = Int16(clamped * 32767.0).littleEndian
+        data.append(Data(bytes: &pcm, count: 2))
+    }
+
+    try! data.write(to: url)
+}
+
+// MARK: - Sampling
+
+func sampleToken(logits: [Float], temperature: Float = 0.8, topK: Int = 50, topP: Float = 0.95) -> Int32 {
+    if temperature <= 0 {
+        return Int32(logits.enumerated().max(by: { $0.element < $1.element })!.offset)
+    }
+
+    var scaled = logits.map { $0 / temperature }
+
+    // Top-k: keep only the top K candidates
+    let indexed = scaled.enumerated().sorted { $0.element > $1.element }
+    let threshold = indexed[min(topK - 1, indexed.count - 1)].element
+    for i in 0..<scaled.count {
+        if scaled[i] < threshold { scaled[i] = -.infinity }
+    }
+
+    // Softmax
+    let maxVal = scaled.max()!
+    var exps = scaled.map { exp($0 - maxVal) }
+    let sum = exps.reduce(0, +)
+    exps = exps.map { $0 / sum }
+
+    // Top-p (nucleus): sort by probability, keep smallest set summing to >= topP
+    let sortedProbs = exps.enumerated().sorted { $0.element > $1.element }
+    var cumProb: Float = 0
+    var allowed = Set<Int>()
+    for (idx, prob) in sortedProbs {
+        cumProb += prob
+        allowed.insert(idx)
+        if cumProb >= topP { break }
+    }
+    // Zero out tokens outside the nucleus
+    for i in 0..<exps.count {
+        if !allowed.contains(i) { exps[i] = 0 }
+    }
+    // Re-normalize
+    let newSum = exps.reduce(0, +)
+    if newSum > 0 { exps = exps.map { $0 / newSum } }
+
+    // Sample
+    let r = Float.random(in: 0..<1)
+    var cumsum: Float = 0
+    for (i, p) in exps.enumerated() {
+        cumsum += p
+        if cumsum >= r { return Int32(i) }
+    }
+    return Int32(exps.count - 1)
+}
+
+// MARK: - Timing
+
+func formatTime(_ seconds: Double) -> String {
+    if seconds < 0.001 { return String(format: "%.2fµs", seconds * 1_000_000) }
+    if seconds < 1.0 { return String(format: "%.1fms", seconds * 1000) }
+    return String(format: "%.2fs", seconds)
+}
+
+func measure<T>(_ label: String, _ block: () throws -> T) rethrows -> T {
+    let start = CFAbsoluteTimeGetCurrent()
+    let result = try block()
+    let elapsed = CFAbsoluteTimeGetCurrent() - start
+    print("  ⏱ \(label): \(formatTime(elapsed))")
+    return result
+}
+
+func measureAsync<T>(_ label: String, _ block: () async throws -> T) async rethrows -> T {
+    let start = CFAbsoluteTimeGetCurrent()
+    let result = try await block()
+    let elapsed = CFAbsoluteTimeGetCurrent() - start
+    print("  ⏱ \(label): \(formatTime(elapsed))")
+    return result
+}
+
+// MARK: - Main pipeline
+
+print("Plapre Pico CoreML TTS Pipeline")
+print("================================\n")
+
+let text = CommandLine.arguments.count > 1 ? CommandLine.arguments[1] : "Hej, mit navn er Daniel."
+let speakerName = CommandLine.arguments.count > 2 ? CommandLine.arguments[2] : "tor"
+let outputPath = CommandLine.arguments.count > 3 ? CommandLine.arguments[3] : "output.wav"
+
+print("Text: \(text)")
+print("Speaker: \(speakerName)")
+print("Output: \(outputPath)\n")
+
+let pipelineStart = CFAbsoluteTimeGetCurrent()
+
+// Load speaker
+let speakerEmb = loadSpeaker(speakerName)
+print("Loaded speaker embedding (\(speakerEmb.count) dims)")
+
+// Tokenize using HuggingFace BPE tokenizer
+let tokenizer = try await measureAsync("Tokenizer load") { try await AutoTokenizer.from(modelFolder: repoRoot) }
+let textTokens = tokenizer.encode(text: text, addSpecialTokens: false).map { Int32($0) }
+print("Tokenized: \(textTokens.count) tokens: \(textTokens)")
+
+// Build input sequence: [placeholder, <text>, tokens..., <audio>]
+var inputSeq: [Int32] = [eosToken, textMarkerToken] + textTokens + [audioMarkerToken]
+let inputLen = inputSeq.count
+print("Input sequence: \(inputLen) tokens")
+
+// Pad to prefillSeqLen
+while inputSeq.count < prefillSeqLen {
+    inputSeq.append(eosToken)
+}
+
+// Load RoPE tables
+print("\nLoading RoPE tables...")
+let ropeCos = loadRopeTable("rope_cos.npy")
+let ropeSin = loadRopeTable("rope_sin.npy")
+print("RoPE cos: \(ropeCos.count) values, sin: \(ropeSin.count) values")
+
+// Compile models
+print("\nCompiling models...")
+var generatedTokens: [Int32] = []
+
+let kanadeModel = try measure("Compile KanadeDecoder") { try compileModel(at: modelURL("KanadeDecoder")) }
+let vocoderModel = try measure("Compile Vocoder") { try compileModel(at: modelURL("Vocoder")) }
+
+if !CommandLine.arguments.contains("--test-audio") {
+let decodeModel = try measure("Compile PlaprePico") { try compileModel(at: modelURL("PlaprePico")) }
+
+// === Step 1: Prefill via decode model (one token at a time) ===
+print("\n--- Prefill (token-by-token, stateless KV cache) ---")
+
+// Allocate KV cache buffers (managed by Swift, passed as model inputs/outputs)
+let numLayers = 30
+let cacheShape = [1, numKvHeads, maxContext, headDim]
+let cacheSize = numKvHeads * maxContext * headDim
+var kvCaches: [MLMultiArray] = []
+for _ in 0..<(numLayers * 2) {
+    kvCaches.append(mlArray([Float](repeating: 0, count: cacheSize), shape: cacheShape))
+}
+
+// Helper to run one token through decode model
+func runDecodeStep(token: Int32, pos: Int, isSpeaker: Bool = false) throws -> [Float] {
+    var maskValues = [Float](repeating: -65504.0, count: maxContext)
+    for j in 0...pos { maskValues[j] = 0.0 }
+
+    let ropeOffset = pos * headDim
+    let cosBuf = Array(ropeCos[ropeOffset..<(ropeOffset + headDim)])
+    let sinBuf = Array(ropeSin[ropeOffset..<(ropeOffset + headDim)])
+
+    var updateMask = [Float](repeating: 0, count: maxContext)
+    updateMask[pos] = 1.0
+
+    var input: [String: MLFeatureValue] = [
+        "input_ids": .init(multiArray: mlArrayInt32([token], shape: [1, 1])),
+        "causal_mask": .init(multiArray: mlArray(maskValues, shape: [1, 1, 1, maxContext])),
+        "cos": .init(multiArray: mlArray(cosBuf, shape: [1, 1, 1, headDim])),
+        "sin": .init(multiArray: mlArray(sinBuf, shape: [1, 1, 1, headDim])),
+        "update_mask": .init(multiArray: mlArray(updateMask, shape: [1, 1, maxContext, 1])),
+        "speaker_embedding": .init(multiArray: mlArray(speakerEmb, shape: [1, speakerDim])),
+        "is_speaker_step": .init(multiArray: mlArray([isSpeaker ? Float(1.0) : Float(0.0)], shape: [1])),
+    ]
+
+    // Add KV cache inputs
+    for i in 0..<numLayers {
+        input["k_cache_\(i)"] = .init(multiArray: kvCaches[2 * i])
+        input["v_cache_\(i)"] = .init(multiArray: kvCaches[2 * i + 1])
+    }
+
+    let provider = try MLDictionaryFeatureProvider(dictionary: input)
+    let output = try decodeModel.prediction(from: provider)
+
+    // Read updated KV caches from output
+    for i in 0..<numLayers {
+        kvCaches[2 * i] = output.featureValue(for: "k_cache_\(i)_out")!.multiArrayValue!
+        kvCaches[2 * i + 1] = output.featureValue(for: "v_cache_\(i)_out")!.multiArrayValue!
+    }
+
+    let logitsArr = output.featureValue(for: "logits")!.multiArrayValue!
+    let count = logitsArr.shape.last!.intValue
+    var result = [Float](repeating: 0, count: count)
+    for i in 0..<count {
+        result[i] = logitsArr[[0, 0, i] as [NSNumber]].floatValue
+    }
+    return result
+}
+
+// The input sequence is: [placeholder(speaker), <text>, tokens..., <audio>]
+// For the speaker token at position 0, we need to handle it differently.
+// The decode model uses embed_tokens, but position 0 should be the speaker projection.
+// WORKAROUND: feed the placeholder token (EOS=2) at position 0. The speaker conditioning
+// won't be perfect since we can't inject the speaker_proj output through the decode model.
+// For proper speaker conditioning, we'd need a dedicated prefill model or a combined model.
+// For now, feed all tokens including the placeholder to validate the pipeline.
+
+let inputTokens: [Int32] = Array(inputSeq.prefix(inputLen))
+print("Processing \(inputTokens.count) input tokens...")
+
+let prefillStart = CFAbsoluteTimeGetCurrent()
+var lastLogits: [Float] = []
+for (i, token) in inputTokens.enumerated() {
+    lastLogits = try runDecodeStep(token: token, pos: i, isSpeaker: i == 0)
+    let argmax = lastLogits.enumerated().max(by: { $0.element < $1.element })!.offset
+    let maxVal = lastLogits.max()!
+    print("  Prefill pos=\(i) token=\(token): argmax=\(argmax) max=\(String(format: "%.4f", maxVal))")
+}
+let prefillElapsed = CFAbsoluteTimeGetCurrent() - prefillStart
+let prefillTokPerSec = Double(inputTokens.count) / prefillElapsed
+print("  ⏱ Prefill: \(formatTime(prefillElapsed)) (\(inputTokens.count) tokens, \(String(format: "%.1f", prefillTokPerSec)) tok/s)")
+
+let firstToken = sampleToken(logits: lastLogits, temperature: 0.8, topK: 50, topP: 0.95)
+print("First generated token: \(firstToken)")
+// Debug
+let dbgLogits = lastLogits
+let sortedIndices = dbgLogits.enumerated().sorted { $0.element > $1.element }
+print("  Top 5: \(sortedIndices.prefix(5).map { "\($0.offset):\($0.element)" })")
+print("  Logits count: \(dbgLogits.count), nonzero: \(dbgLogits.filter { $0 != 0 }.count)")
+print("  Speaker emb first 3: \(speakerEmb.prefix(3))")
+
+// === Step 2: Autoregressive decode ===
+print("\n--- Decode ---")
+generatedTokens = [firstToken]
+let maxTokens = 500
+
+print("Generating up to \(maxTokens) tokens...")
+let decodeStart = CFAbsoluteTimeGetCurrent()
+var nextToken = firstToken
+var consecutiveNonAudio = 0
+let nonAudioStopThreshold = 10  // stop after this many consecutive non-audio tokens
+for step in 1..<maxTokens {
+    let pos = inputLen + step - 1
+
+    let logits = try runDecodeStep(token: nextToken, pos: pos)
+    nextToken = sampleToken(logits: logits, temperature: 0.8, topK: 50, topP: 0.95)
+    generatedTokens.append(nextToken)
+
+    if nextToken == eosToken {
+        print("  EOS at step \(step)")
+        break
+    }
+
+    // Track consecutive non-audio tokens — model may be done speaking
+    if nextToken >= audioTokenOffset && nextToken <= 20801 {
+        consecutiveNonAudio = 0
+    } else {
+        consecutiveNonAudio += 1
+        if consecutiveNonAudio >= nonAudioStopThreshold {
+            print("  Stopping: \(nonAudioStopThreshold) consecutive non-audio tokens at step \(step)")
+            break
+        }
+    }
+
+    if step % 25 == 0 {
+        let elapsed = CFAbsoluteTimeGetCurrent() - decodeStart
+        let tokPerSec = Double(step) / elapsed
+        print("  Step \(step) (\(Float(step) / 25.0)s audio) — \(formatTime(elapsed)) elapsed, \(String(format: "%.1f", tokPerSec)) tok/s")
+    }
+}
+let decodeElapsed = CFAbsoluteTimeGetCurrent() - decodeStart
+let decodeSteps = generatedTokens.count - 1  // first token came from prefill
+let decodeTokPerSec = Double(decodeSteps) / decodeElapsed
+let audioSeconds = Float(generatedTokens.filter { $0 >= audioTokenOffset && $0 <= 20801 }.count) / 25.0
+let rtf = Float(decodeElapsed) / audioSeconds  // real-time factor: wall time / audio duration
+print("  ⏱ Decode: \(formatTime(decodeElapsed)) (\(decodeSteps) steps, \(String(format: "%.1f", decodeTokPerSec)) tok/s)")
+print("  ⏱ Audio generated: \(String(format: "%.1f", audioSeconds))s — RTF \(String(format: "%.2f", rtf))x (1.0 = realtime)")
+
+} // end if !test-audio
+
+var audioTokens: [Int32]
+let testAudioOnly = CommandLine.arguments.contains("--test-audio")
+
+if testAudioOnly {
+    // Skip LLM, use known-good tokens from Python pipeline for audio testing
+    print("\n--- Using hardcoded test tokens (--test-audio) ---")
+    audioTokens = [11620, 17958, 13738, 15707, 12635, 12635, 12131, 12637, 20677, 12903,
+                   17769, 17841, 20016, 20080, 17520, 20080, 17528, 14832, 14774, 12200,
+                   12199, 12263, 11693, 11622, 12130, 12066, 12050, 12050, 12050, 12050,
+                   14578, 14642, 14610, 14082, 12058, 11482, 11474, 14538, 14610, 14642,
+                   14610, 14082, 14082, 11490, 11482, 11482, 11482, 11482, 11482, 11474,
+                   11410, 11394, 12066, 12058, 14610, 14610, 14098, 11490, 11482, 11490,
+                   11482, 11482, 11482, 11482, 11482, 11474, 11410, 11394, 11394, 11954,
+                   12010, 12002, 11426, 11418, 11026, 14618, 14082, 12061, 19682, 19933,
+                   20590, 19877, 17770, 17322, 14832, 14760, 12192, 12200, 12192, 12200,
+                   12199, 12263, 11693, 11686, 11677, 11686, 8914, 8978, 8914, 8978]
+    generatedTokens = audioTokens
+} else {
+    audioTokens = generatedTokens.filter { $0 >= audioTokenOffset && $0 <= 20801 }
+}
+
+print("\nGenerated \(generatedTokens.count) tokens, \(audioTokens.count) audio (\(Float(audioTokens.count) / 25.0)s)")
+print("All tokens: \(generatedTokens.map { String($0) }.joined(separator: ", "))")
+print("Audio tokens: \(audioTokens.prefix(20).map { String($0) }.joined(separator: ", "))...")
+
+if audioTokens.isEmpty {
+    print("No audio tokens generated!")
+    exit(1)
+}
+
+// === Step 3: Kanade + Vocoder in chunks ===
+// Kanade expects exactly 100 tokens (4s at 25 tokens/sec).
+// Process audio tokens in 100-token chunks, concatenate waveforms.
+let kanadeChunkSize = 100
+let numChunks = (audioTokens.count + kanadeChunkSize - 1) / kanadeChunkSize
+print("\n--- Kanade + Vocoder (\(numChunks) chunk\(numChunks == 1 ? "" : "s") of \(kanadeChunkSize) tokens) ---")
+
+var waveform: [Float] = []
+let audioDecodeStart = CFAbsoluteTimeGetCurrent()
+
+for chunkIdx in 0..<numChunks {
+    let chunkStart = CFAbsoluteTimeGetCurrent()
+    let start = chunkIdx * kanadeChunkSize
+    let end = min(start + kanadeChunkSize, audioTokens.count)
+    let chunkTokens = Array(audioTokens[start..<end])
+
+    // Convert to Kanade indices (subtract audio offset) and pad to chunk size
+    var kanadeIndices = chunkTokens.map { $0 - Int32(audioTokenOffset) }
+    let actualCount = kanadeIndices.count
+    while kanadeIndices.count < kanadeChunkSize {
+        kanadeIndices.append(kanadeIndices.last ?? 0)  // repeat last token as padding
+    }
+
+    // Kanade: tokens → mel
+    let kanadeStart = CFAbsoluteTimeGetCurrent()
+    let kanadeInput: [String: MLFeatureValue] = [
+        "token_indices": .init(multiArray: mlArrayInt32(kanadeIndices, shape: [kanadeChunkSize])),
+        "speaker_embedding": .init(multiArray: mlArrayFloat32(speakerEmb, shape: [1, speakerDim])),
+    ]
+    let kanadeProvider = try MLDictionaryFeatureProvider(dictionary: kanadeInput)
+    let kanadeOutput = try kanadeModel.prediction(from: kanadeProvider)
+    let mel = kanadeOutput.featureValue(for: "mel")!.multiArrayValue!
+    let kanadeElapsed = CFAbsoluteTimeGetCurrent() - kanadeStart
+
+    // Vocoder: mel → waveform
+    let vocoderStart = CFAbsoluteTimeGetCurrent()
+    let vocoderInput: [String: MLFeatureValue] = [
+        "mel": .init(multiArray: mel),
+    ]
+    let vocoderProvider = try MLDictionaryFeatureProvider(dictionary: vocoderInput)
+    let vocoderOutput = try vocoderModel.prediction(from: vocoderProvider)
+    let chunkWaveform = readFloat32Array(vocoderOutput.featureValue(for: "waveform")!.multiArrayValue!)
+    let vocoderElapsed = CFAbsoluteTimeGetCurrent() - vocoderStart
+
+    // If this chunk was padded, trim the waveform proportionally
+    let samplesPerToken = chunkWaveform.count / kanadeChunkSize  // 960 samples per token at 24kHz
+    let usableSamples = actualCount * samplesPerToken
+    waveform.append(contentsOf: chunkWaveform.prefix(usableSamples))
+
+    let chunkElapsed = CFAbsoluteTimeGetCurrent() - chunkStart
+    let chunkDuration = String(format: "%.1f", Float(usableSamples) / Float(sampleRate))
+    print("  Chunk \(chunkIdx + 1)/\(numChunks): \(actualCount) tokens → \(chunkDuration)s audio — Kanade \(formatTime(kanadeElapsed)), Vocoder \(formatTime(vocoderElapsed)), total \(formatTime(chunkElapsed))")
+}
+let audioDecodeElapsed = CFAbsoluteTimeGetCurrent() - audioDecodeStart
+print("  ⏱ Audio decode total: \(formatTime(audioDecodeElapsed)) (\(numChunks) chunk\(numChunks == 1 ? "" : "s"))")
+
+print("Total waveform: \(waveform.count) samples (\(String(format: "%.1f", Float(waveform.count) / Float(sampleRate)))s)")
+
+// === Write WAV ===
+let outputURL = URL(fileURLWithPath: outputPath)
+writeWAV(waveform, to: outputURL)
+print("\nSaved to \(outputPath)")
+
+// === Timing Summary ===
+let pipelineElapsed = CFAbsoluteTimeGetCurrent() - pipelineStart
+let totalAudioDuration = Float(waveform.count) / Float(sampleRate)
+print("\n========== Timing Summary ==========")
+print("  Total pipeline:    \(formatTime(pipelineElapsed))")
+print("  Audio output:      \(String(format: "%.1f", totalAudioDuration))s")
+print("  Overall RTF:       \(String(format: "%.2f", Float(pipelineElapsed) / totalAudioDuration))x")
+print("====================================")
+print("Done!")