Initial commit: E4B-MarkBase model integration with passing tests

- E4B-MarkBase model (42 layers, 4.4GB) loaded successfully - All Phase 1-6 tests passed (model loading, forward pass, vision/audio towers, token generation, performance) - All stress tests passed (5/5 in 127.6s) - Concurrent inference - Memory stress (67.5 tok/s, 0 NaN) - Continuous generation - Batch processing - Long-running stability - Swift Metal inference engine with multimodal support
2026-06-23 18:12:35 +08:00
commit ac75faa0cc
301 changed files with 63426 additions and 0 deletions
@@ -0,0 +1,230 @@
+import AVFoundation
+import Metal
+
+public final class AudioFeatureExtractor {
+    public let sampleRate: Int
+    public let nMels: Int
+    public let nFft: Int
+    public let hopLength: Int
+    public let fMin: Float
+    public let fMax: Float
+    
+    public init(
+        sampleRate: Int = 16000,
+        nMels: Int = 128,
+        nFft: Int = 400,
+        hopLength: Int = 160,
+        fMin: Float = 0,
+        fMax: Float = 8000
+    ) {
+        self.sampleRate = sampleRate
+        self.nMels = nMels
+        self.nFft = nFft
+        self.hopLength = hopLength
+        self.fMin = fMin
+        self.fMax = fMax
+    }
+    
+    public func extractMelSpectrogram(from audioData: [Float]) -> [[Float]] {
+        let numFrames = (audioData.count - nFft) / hopLength + 1
+        var melSpec = [[Float]](repeating: [Float](repeating: 0, count: nMels), count: numFrames)
+        
+        for frameIdx in 0..<numFrames {
+            let startIdx = frameIdx * hopLength
+            let endIdx = min(startIdx + nFft, audioData.count)
+            
+            // Zero-pad if frame is shorter than nFft
+            var frame = [Float](repeating: 0, count: nFft)
+            for i in startIdx..<endIdx {
+                frame[i - startIdx] = audioData[i]
+            }
+            
+            let windowedFrame = applyHannWindow(frame)
+            let spectrum = computeSpectrum(windowedFrame)
+            let melEnergies = computeMelEnergies(spectrum)
+            
+            melSpec[frameIdx] = melEnergies
+        }
+        
+        return melSpec
+    }
+    
+    private func applyHannWindow(_ frame: [Float]) -> [Float] {
+        let n = frame.count
+        return frame.enumerated().map { i, val in
+            val * 0.5 * (1.0 - cos(2.0 * Float.pi * Float(i) / Float(n - 1)))
+        }
+    }
+    
+    private func computeSpectrum(_ frame: [Float]) -> [Float] {
+        let n = frame.count
+        var spectrum = [Float](repeating: 0, count: n / 2 + 1)
+        
+        for k in 0..<spectrum.count {
+            var real: Float = 0
+            var imag: Float = 0
+            for i in 0..<n {
+                let angle = -2.0 * Float.pi * Float(k) * Float(i) / Float(n)
+                real += frame[i] * cos(angle)
+                imag += frame[i] * sin(angle)
+            }
+            spectrum[k] = sqrt(real * real + imag * imag)
+        }
+        
+        return spectrum
+    }
+    
+    private func computeMelEnergies(_ spectrum: [Float]) -> [Float] {
+        var melEnergies = [Float](repeating: 0, count: nMels)
+        
+        let melPoints = createMelFilterbank()
+        
+        for melIdx in 0..<nMels {
+            var sum: Float = 0
+            for fftIdx in 0..<spectrum.count {
+                sum += spectrum[fftIdx] * melPoints[melIdx][fftIdx]
+            }
+            melEnergies[melIdx] = log10(max(sum, 1e-10))
+        }
+        
+        return melEnergies
+    }
+    
+    private func createMelFilterbank() -> [[Float]] {
+        var filterbank = [[Float]](repeating: [Float](repeating: 0, count: nFft / 2 + 1), count: nMels)
+        
+        let melMin = hzToMel(fMin)
+        let melMax = hzToMel(fMax)
+        
+        let melPoints = (0..<nMels + 2).map { i in
+            melMin + Float(i) * (melMax - melMin) / Float(nMels + 1)
+        }
+        
+        let hzPoints = melPoints.map { melToHz($0) }
+        let binPoints = hzPoints.map { Int(round($0 * Float(nFft) / Float(sampleRate))) }
+        
+        for melIdx in 0..<nMels {
+            let leftBin = binPoints[melIdx]
+            let centerBin = binPoints[melIdx + 1]
+            let rightBin = binPoints[melIdx + 2]
+            
+            for bin in leftBin..<centerBin {
+                if bin < filterbank[melIdx].count {
+                    filterbank[melIdx][bin] = Float(bin - leftBin) / Float(centerBin - leftBin)
+                }
+            }
+            for bin in centerBin..<rightBin {
+                if bin < filterbank[melIdx].count {
+                    filterbank[melIdx][bin] = Float(rightBin - bin) / Float(rightBin - centerBin)
+                }
+            }
+        }
+        
+        return filterbank
+    }
+    
+    private func hzToMel(_ hz: Float) -> Float {
+        2595.0 * log10(1.0 + hz / 700.0)
+    }
+    
+    private func melToHz(_ mel: Float) -> Float {
+        700.0 * (pow(10.0, mel / 2595.0) - 1.0)
+    }
+    
+    // ── GPU-accelerated mel spectrogram ──
+    
+    public func extractMelSpectrogramGPU(
+        engine: MarkBaseEngine,
+        audioData: [Float]
+    ) throws -> [[Float]] {
+        let device = engine.device
+        let spectrumSize = nFft / 2 + 1
+        let numFrames = (audioData.count - nFft) / hopLength + 1
+        let melBufferSize = numFrames * nMels
+
+        let filterbank2D = createMelFilterbank()
+        var flatFilterbank = [Float](repeating: 0, count: nMels * spectrumSize)
+        for m in 0..<nMels {
+            for b in 0..<spectrumSize {
+                flatFilterbank[m * spectrumSize + b] = filterbank2D[m][b]
+            }
+        }
+
+        let cmdBuf = engine.commandQueue.makeCommandBuffer()!
+        let audioBuf = device.makeBuffer(bytes: audioData, length: audioData.count * 4)!
+        let filterbankBuf = device.makeBuffer(bytes: flatFilterbank, length: flatFilterbank.count * 4)!
+        let spectrumBuf = device.makeBuffer(length: numFrames * spectrumSize * 4)!
+        let melBuf = device.makeBuffer(length: melBufferSize * 4)!
+
+        let psoDFT = try engine.pipeline(named: "audio_dft_magnitude")
+        let psoMel = try engine.pipeline(named: "audio_mel_filterbank")
+
+        do {
+            let enc = cmdBuf.makeComputeCommandEncoder()!
+            enc.setComputePipelineState(psoDFT)
+            enc.setBuffer(audioBuf, offset: 0, index: 0)
+            enc.setBuffer(spectrumBuf, offset: 0, index: 1)
+            var n = UInt32(nFft); enc.setBytes(&n, length: 4, index: 2)
+            var h = UInt32(hopLength); enc.setBytes(&h, length: 4, index: 3)
+            var nf = UInt32(numFrames); enc.setBytes(&nf, length: 4, index: 4)
+            var ss = UInt32(spectrumSize); enc.setBytes(&ss, length: 4, index: 5)
+            var al = UInt32(audioData.count); enc.setBytes(&al, length: 4, index: 6)
+            let grid = MTLSize(width: numFrames, height: spectrumSize, depth: 1)
+            let tg = MTLSize(width: 8, height: 8, depth: 1)
+            enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
+            enc.endEncoding()
+        }
+        do {
+            let enc = cmdBuf.makeComputeCommandEncoder()!
+            enc.setComputePipelineState(psoMel)
+            enc.setBuffer(spectrumBuf, offset: 0, index: 0)
+            enc.setBuffer(filterbankBuf, offset: 0, index: 1)
+            enc.setBuffer(melBuf, offset: 0, index: 2)
+            var ss = UInt32(spectrumSize); enc.setBytes(&ss, length: 4, index: 3)
+            var nm = UInt32(nMels); enc.setBytes(&nm, length: 4, index: 4)
+            var nf = UInt32(numFrames); enc.setBytes(&nf, length: 4, index: 5)
+            let grid = MTLSize(width: numFrames, height: nMels, depth: 1)
+            let tg = MTLSize(width: 8, height: 8, depth: 1)
+            enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
+            enc.endEncoding()
+        }
+
+        cmdBuf.commit()
+        cmdBuf.waitUntilCompleted()
+
+        let ptr = melBuf.contents().assumingMemoryBound(to: Float.self)
+        let flat = Array(UnsafeBufferPointer(start: ptr, count: melBufferSize))
+        var result = [[Float]](repeating: [Float](repeating: 0, count: nMels), count: numFrames)
+        for f in 0..<numFrames {
+            for m in 0..<nMels {
+                result[f][m] = flat[f * nMels + m]
+            }
+        }
+        return result
+    }
+
+    public func loadAudioFile(url: URL) throws -> [Float] {
+        let asset = AVURLAsset(url: url)
+        let reader = try AVAssetReader(asset: asset)
+        
+        let output = AVAssetReaderAudioMixOutput(audioTracks: asset.tracks, audioSettings: nil)
+        reader.add(output)
+        reader.startReading()
+        
+        var samples: [Float] = []
+        while reader.status == .reading {
+            let buffer = output.copyNextSampleBuffer()
+            if let buffer = buffer {
+                let blockBuffer = CMSampleBufferGetDataBuffer(buffer)
+                if let blockBuffer = blockBuffer {
+                    let length = CMBlockBufferGetDataLength(blockBuffer)
+                    var data = [Float](repeating: 0, count: length / MemoryLayout<Float>.stride)
+                    CMBlockBufferCopyDataBytes(blockBuffer, atOffset: 0, dataLength: length, destination: &data)
+                    samples.append(contentsOf: data)
+                }
+            }
+        }
+        
+        return samples
+    }
+}