Add DeepFilterNet speech enhancement model

.github/workflows/tests.yaml

@@ -23,6 +23,8 @@ jobs:
             suite: MLXAudioTests/MLXAudioSTTTests
           - name: VAD
             suite: MLXAudioTests/MLXAudioVADTests
+          - name: STS
+            suite: MLXAudioTests/MLXAudioSTSTests
 
     steps:
       - uses: actions/checkout@v4

.gitignore

@@ -1,5 +1,7 @@
 .build
 .swiftpm
+.scratch
+.xcodebuild
 .DS_Store
 .xcbkptlist
 .claude
@@ -13,3 +15,6 @@ xcuserdata/*
 *.xcuserdatad
 *.xcuserdatad/xcdebugger
 *.xcuserdatad/xcdebugger/Breakpoints_v2.xcbkptlist
+.xcode_derived/
+default.profraw
+scripts/

README.md

@@ -17,7 +17,7 @@ MLXAudio follows a modular design allowing you to import only what you need:
 - **MLXAudioTTS**: Text-to-Speech models (Qwen3-TTS, Fish Audio S2 Pro, Soprano, VyvoTTS, Orpheus, Marvis TTS, Pocket TTS)
 - **MLXAudioSTT**: Speech-to-Text models (Qwen3-ASR, Voxtral Realtime, Parakeet, GLMASR)
 - **MLXAudioVAD**: Voice Activity Detection & Speaker Diarization (Sortformer, SmartTurn)
-- **MLXAudioSTS**: Speech-to-Speech models (LFM2.5-Audio, SAM-Audio, MossFormer2-SE)
+- **MLXAudioSTS**: Speech-to-Speech models (LFM2.5-Audio, SAM-Audio, MossFormer2-SE, DeepFilterNet)
 - **MLXAudioUI**: SwiftUI components for audio interfaces
 
 ## Installation
@@ -143,6 +143,7 @@ for try await event in model.generateStream(text: text, parameters: parameters)
 | LFM2.5-Audio | [LFM Audio README](Sources/MLXAudioSTS/Models/LFMAudio/README.md) | [mlx-community/LFM2.5-Audio-1.5B-6bit](https://huggingface.co/mlx-community/LFM2.5-Audio-1.5B-6bit) |
 | SAM-Audio | [SAM Audio README](Sources/MLXAudioSTS/Models/SAMAudio/README.md) | [mlx-community/sam-audio-large-fp16](https://huggingface.co/mlx-community/sam-audio-large-fp16) |
 | MossFormer2-SE | — | [starkdmi/MossFormer2-SE-fp16](https://huggingface.co/starkdmi/MossFormer2-SE-fp16) |
+| DeepFilterNet | [DeepFilterNet README](Sources/MLXAudioSTS/Models/DeepFilterNet/README.md) | [mlx-community/DeepFilterNet-mlx](https://huggingface.co/mlx-community/DeepFilterNet-mlx) |
 
 ### VAD / Speaker Diarization Models
 

Sources/MLXAudioCore/AudioUtils.swift

@@ -233,6 +233,8 @@ public class StreamingWAVWriter {
     private let sampleRate: Double
     private var audioFile: AVAudioFile?
     private let format: AVAudioFormat
+    private var reusableBuffer: AVAudioPCMBuffer?
+    private var reusableCapacity: Int = 0
     public private(set) var framesWritten: Int = 0
 
     public init(url: URL, sampleRate: Double) throws {
@@ -252,39 +254,95 @@ public class StreamingWAVWriter {
 
     /// Write a chunk of audio samples to the file.
     public func writeChunk(_ samples: [Float]) throws {
+        try samples.withUnsafeBytes { raw in
+            let data = Data(raw)
+            try writeChunkData(data, sampleCount: samples.count)
+        }
+    }
+
+    /// Write a chunk of audio samples from an MLX array directly.
+    /// Expects a 1D float32 tensor.
+    public func writeChunk(_ samples: MLXArray) throws {
+        let f32 = samples.dtype == .float32 ? samples : samples.asType(.float32)
+        let sampleCount = f32.size
+        guard sampleCount > 0 else { return }
+        let data = f32.asData(access: .noCopyIfContiguous).data
+        try writeChunkData(data, sampleCount: sampleCount)
+    }
+
+    /// Write a chunk of float32 samples from a byte buffer.
+    public func writeChunkData(_ samplesData: Data, sampleCount: Int) throws {
         guard let audioFile else {
             throw NSError(
                 domain: "StreamingWAVWriter",
                 code: 2,
                 userInfo: [NSLocalizedDescriptionKey: "Audio file not initialized or already finalized"]
             )
         }
+        guard sampleCount >= 0 else {
+            throw NSError(
+                domain: "StreamingWAVWriter",
+                code: 4,
+                userInfo: [NSLocalizedDescriptionKey: "Invalid negative sample count: \(sampleCount)"]
+            )
+        }
+        guard sampleCount > 0 else { return }
 
-        let frameCount = AVAudioFrameCount(samples.count)
-        guard let buffer = AVAudioPCMBuffer(pcmFormat: format, frameCapacity: frameCount) else {
+        let expectedBytes = sampleCount * MemoryLayout<Float>.size
+        guard samplesData.count >= expectedBytes else {
             throw NSError(
                 domain: "StreamingWAVWriter",
-                code: 3,
-                userInfo: [NSLocalizedDescriptionKey: "Failed to create audio buffer"]
+                code: 5,
+                userInfo: [NSLocalizedDescriptionKey: "Insufficient sample data bytes (\(samplesData.count) < \(expectedBytes))"]
             )
         }
 
+        let frameCount = AVAudioFrameCount(sampleCount)
+        let buffer = try getReusableBuffer(minFrameCapacity: sampleCount)
         buffer.frameLength = frameCount
 
-        if let channelData = buffer.floatChannelData {
-            for i in 0 ..< samples.count {
-                channelData[0][i] = samples[i]
+        guard let channelData = buffer.floatChannelData else {
+            throw NSError(
+                domain: "StreamingWAVWriter",
+                code: 6,
+                userInfo: [NSLocalizedDescriptionKey: "Failed to access output channel data"]
+            )
+        }
+        samplesData.withUnsafeBytes { src in
+            if let srcBase = src.baseAddress {
+                memcpy(channelData[0], srcBase, expectedBytes)
             }
         }
 
         try audioFile.write(from: buffer)
-        framesWritten += samples.count
+        framesWritten += sampleCount
     }
 
     /// Finalize the WAV file and return the URL.
     /// After calling this method, no more chunks can be written.
     public func finalize() -> URL {
         audioFile = nil // Close the file by releasing the reference
+        reusableBuffer = nil
+        reusableCapacity = 0
         return url
     }
+
+    private func getReusableBuffer(minFrameCapacity: Int) throws -> AVAudioPCMBuffer {
+        if let reusableBuffer, reusableCapacity >= minFrameCapacity {
+            return reusableBuffer
+        }
+        guard let newBuffer = AVAudioPCMBuffer(
+            pcmFormat: format,
+            frameCapacity: AVAudioFrameCount(minFrameCapacity)
+        ) else {
+            throw NSError(
+                domain: "StreamingWAVWriter",
+                code: 3,
+                userInfo: [NSLocalizedDescriptionKey: "Failed to create audio buffer"]
+            )
+        }
+        reusableBuffer = newBuffer
+        reusableCapacity = minFrameCapacity
+        return newBuffer
+    }
 }

Sources/MLXAudioCore/ModelUtils.swift

@@ -16,7 +16,10 @@ public enum ModelUtils {
         )
         let configJSON = try JSONSerialization.jsonObject(with: Data(contentsOf: modelURL.appendingPathComponent("config.json")))
         if let config = configJSON as? [String: Any] {
-            return (config["model_type"] as? String) ?? (config["architecture"] as? String) ?? modelNameComponents?.first?.lowercased()
+            return (config["model_type"] as? String)
+                ?? (config["architecture"] as? String)
+                ?? (config["model_version"] as? String)
+                ?? modelNameComponents?.first?.lowercased()
         }
         return nil
     }

Sources/MLXAudioSTS/Models/DeepFilterNet/DeepFilterNetConfig.swift

@@ -0,0 +1,125 @@
+import Foundation
+
+/// Model configuration for DeepFilterNet, decoded from `config.json` in the model directory.
+public struct DeepFilterNetConfig: Codable, Sendable {
+    public var sampleRate: Int = 48_000
+    public var fftSize: Int = 960
+    public var hopSize: Int = 480
+    public var minNbErbFreqs: Int = 2
+
+    public var nbErb: Int = 32
+    public var nbDf: Int = 96
+    public var dfOrder: Int = 5
+    public var dfLookahead: Int = 2
+    public var convLookahead: Int = 2
+
+    public var convCh: Int = 64
+    public var convKEnc: Int = 1
+    public var convKDec: Int = 1
+    public var convWidthFactor: Int = 1
+    public var embHiddenDim: Int = 256
+    public var embNumLayers: Int = 3
+    public var dfHiddenDim: Int = 256
+    public var dfNumLayers: Int = 2
+    public var gruGroups: Int = 8
+    public var linearGroups: Int = 16
+    public var encLinearGroups: Int = 32
+    public var groupShuffle: Bool = false
+    public var encConcat: Bool = false
+    public var dfGruSkip: String = "groupedlinear"
+
+    public var convKernel: [Int] = [1, 3]
+    public var convtKernel: [Int] = [1, 3]
+    public var convKernelInp: [Int] = [3, 3]
+    public var dfPathwayKernelSizeT: Int = 5
+
+    public var lsnrMax: Int = 35
+    public var lsnrMin: Int = -15
+    public var modelVersion: String = "DeepFilterNet3"
+    public var erbWidths: [Int]? = nil
+
+    public var freqBins: Int { fftSize / 2 + 1 }
+
+    public var modelType: String {
+        switch modelVersion.lowercased() {
+        case "deepfilternet":
+            return "deepfilternet"
+        case "deepfilternet2":
+            return "deepfilternet2"
+        default:
+            return "deepfilternet3"
+        }
+    }
+
+    enum CodingKeys: String, CodingKey {
+        case sampleRate
+        case fftSize
+        case hopSize
+        case minNbErbFreqs
+        case nbErb
+        case nbDf
+        case dfOrder
+        case dfLookahead
+        case convLookahead
+        case convCh
+        case convKEnc
+        case convKDec
+        case convWidthFactor
+        case embHiddenDim
+        case embNumLayers
+        case dfHiddenDim
+        case dfNumLayers
+        case gruGroups
+        case linearGroups
+        case encLinearGroups
+        case groupShuffle
+        case encConcat
+        case dfGruSkip
+        case convKernel
+        case convtKernel
+        case convKernelInp
+        case dfPathwayKernelSizeT
+        case lsnrMax
+        case lsnrMin
+        case modelVersion
+        case erbWidths
+    }
+
+    public init() {}
+
+    public init(from decoder: Decoder) throws {
+        let c = try decoder.container(keyedBy: CodingKeys.self)
+        sampleRate = try c.decodeIfPresent(Int.self, forKey: .sampleRate) ?? sampleRate
+        fftSize = try c.decodeIfPresent(Int.self, forKey: .fftSize) ?? fftSize
+        hopSize = try c.decodeIfPresent(Int.self, forKey: .hopSize) ?? hopSize
+        minNbErbFreqs = try c.decodeIfPresent(Int.self, forKey: .minNbErbFreqs) ?? minNbErbFreqs
+        nbErb = try c.decodeIfPresent(Int.self, forKey: .nbErb) ?? nbErb
+        nbDf = try c.decodeIfPresent(Int.self, forKey: .nbDf) ?? nbDf
+        dfOrder = try c.decodeIfPresent(Int.self, forKey: .dfOrder) ?? dfOrder
+        dfLookahead = try c.decodeIfPresent(Int.self, forKey: .dfLookahead) ?? dfLookahead
+        convLookahead = try c.decodeIfPresent(Int.self, forKey: .convLookahead) ?? convLookahead
+        convCh = try c.decodeIfPresent(Int.self, forKey: .convCh) ?? convCh
+        convKEnc = try c.decodeIfPresent(Int.self, forKey: .convKEnc) ?? convKEnc
+        convKDec = try c.decodeIfPresent(Int.self, forKey: .convKDec) ?? convKDec
+        convWidthFactor = try c.decodeIfPresent(Int.self, forKey: .convWidthFactor) ?? convWidthFactor
+        embHiddenDim = try c.decodeIfPresent(Int.self, forKey: .embHiddenDim) ?? embHiddenDim
+        embNumLayers = try c.decodeIfPresent(Int.self, forKey: .embNumLayers) ?? embNumLayers
+        dfHiddenDim = try c.decodeIfPresent(Int.self, forKey: .dfHiddenDim) ?? dfHiddenDim
+        dfNumLayers = try c.decodeIfPresent(Int.self, forKey: .dfNumLayers) ?? dfNumLayers
+        gruGroups = try c.decodeIfPresent(Int.self, forKey: .gruGroups) ?? gruGroups
+        linearGroups = try c.decodeIfPresent(Int.self, forKey: .linearGroups) ?? linearGroups
+        encLinearGroups = try c.decodeIfPresent(Int.self, forKey: .encLinearGroups) ?? encLinearGroups
+        groupShuffle = try c.decodeIfPresent(Bool.self, forKey: .groupShuffle) ?? groupShuffle
+        encConcat = try c.decodeIfPresent(Bool.self, forKey: .encConcat) ?? encConcat
+        dfGruSkip = try c.decodeIfPresent(String.self, forKey: .dfGruSkip) ?? dfGruSkip
+        convKernel = try c.decodeIfPresent([Int].self, forKey: .convKernel) ?? convKernel
+        convtKernel = try c.decodeIfPresent([Int].self, forKey: .convtKernel) ?? convtKernel
+        convKernelInp = try c.decodeIfPresent([Int].self, forKey: .convKernelInp) ?? convKernelInp
+        dfPathwayKernelSizeT =
+            try c.decodeIfPresent(Int.self, forKey: .dfPathwayKernelSizeT) ?? dfPathwayKernelSizeT
+        lsnrMax = try c.decodeIfPresent(Int.self, forKey: .lsnrMax) ?? lsnrMax
+        lsnrMin = try c.decodeIfPresent(Int.self, forKey: .lsnrMin) ?? lsnrMin
+        modelVersion = try c.decodeIfPresent(String.self, forKey: .modelVersion) ?? modelVersion
+        erbWidths = try c.decodeIfPresent([Int].self, forKey: .erbWidths) ?? erbWidths
+    }
+}