Skip to content

Fix Kokoro v2 source_noise dtype and distribution #447

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Alex-Wengg merged 1 commit into from
Mar 28, 2026
+14 −23
Merged

Fix Kokoro v2 source_noise dtype and distribution #447

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 1 addition & 11 deletions Sources/FluidAudio/ModelNames.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -535,23 +535,13 @@ public enum ModelNames {

/// Underlying model bundle filename.
public var fileName: String {
#if os(iOS)
// Use v1 models on iOS - v2 fp16 models cause warm-up hangs
// Use v1 models on all platforms - v2 has source_noise issues
switch self {
case .fiveSecond:
return "kokoro_21_5s.mlmodelc"
case .fifteenSecond:
return "kokoro_21_15s.mlmodelc"
}
#else
// Use v2 models on macOS - fp16 ANE optimization (1.67x faster)
switch self {
case .fiveSecond:
return "kokoro_21_5s_v2.mlmodelc"
case .fifteenSecond:
return "kokoro_21_15s_v2.mlmodelc"
}
#endif
}

/// Approximate maximum duration in seconds handled by the variant.
Expand Down
25 changes: 13 additions & 12 deletions Sources/FluidAudio/TTS/Kokoro/Pipeline/Synthesize/KokoroSynthesizer.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -411,20 +411,21 @@ public struct KokoroSynthesizer {
throw TTSError.processingFailed("Failed to extract 'audio' output. Features: \(names)")
}

// Optional: trim to audio_length_samples if provided
// Compute audio length from pred_dur (model's audio_length_samples output is broken)
var effectiveCount = audioArrayUnwrapped.count
if let lenFV = output.featureValue(for: "audio_length_samples") {
var n: Int = 0
if let lenArray = lenFV.multiArrayValue, lenArray.count > 0 {
n = lenArray[0].intValue
} else if lenFV.type == .int64 {
n = Int(lenFV.int64Value)
} else if lenFV.type == .double {
n = Int(lenFV.doubleValue)

if let predDurArray = output.featureValue(for: "pred_dur")?.multiArrayValue {
// Sum pred_dur to get total frames
var totalFrames: Float = 0.0
let predDurPtr = predDurArray.dataPointer.bindMemory(to: Float.self, capacity: predDurArray.count)
for i in 0..<predDurArray.count {
totalFrames += predDurPtr[i]
}
Comment on lines +419 to 423
Copy link
Copy Markdown
Contributor

@devin-ai-integration devin-ai-integration bot Mar 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔴 pred_dur dataPointer bound as Float.self without checking actual MLMultiArray data type

The new pred_dur reading code at KokoroSynthesizer.swift:420 assumes the array is float32 by using bindMemory(to: Float.self), but never checks predDurArray.dataType. If the CoreML model outputs pred_dur in float16 format (2 bytes/element), binding as Float.self (4 bytes/element) causes an out-of-bounds memory read (reading count * 4 bytes from a count * 2 byte buffer) and produces garbage frame counts, leading to incorrect audio trimming. This is inconsistent with the audio output handling just below at KokoroSynthesizer.swift:448, which correctly checks audioArrayUnwrapped.dataType == .float32 before pointer-binding and has a safe fallback using [i].floatValue.

Safe pattern used for audio (line 448) but not for pred_dur 
// Audio output: correctly checks dataType
if audioArrayUnwrapped.dataType == .float32 {
    let sourcePointer = audioArrayUnwrapped.dataPointer.bindMemory(...)
} else {
    // safe fallback via NSNumber
}

// pred_dur: no dataType check
let predDurPtr = predDurArray.dataPointer.bindMemory(to: Float.self, ...)
Suggested change
var totalFrames: Float = 0.0
let predDurPtr = predDurArray.dataPointer.bindMemory(to: Float.self, capacity: predDurArray.count)
for i in 0..<predDurArray.count {
totalFrames += predDurPtr[i]
}
var totalFrames: Float = 0.0
if predDurArray.dataType == .float32 {
let predDurPtr = predDurArray.dataPointer.bindMemory(to: Float.self, capacity: predDurArray.count)
for i in 0..<predDurArray.count {
totalFrames += predDurPtr[i]
}
} else {
for i in 0..<predDurArray.count {
totalFrames += predDurArray[i].floatValue
}
}
Open in Devin Review

Was this helpful? React with 👍 or 👎 to provide feedback.

n = max(0, n)
if n > 0 && n <= audioArrayUnwrapped.count {
effectiveCount = n

// Convert frames to samples: frames * 600 samples/frame
let predictedSamples = Int(round(totalFrames * 600.0))
if predictedSamples > 0 {
effectiveCount = min(predictedSamples, audioArrayUnwrapped.count)
}
}

Expand Down
Loading