react-native-nitro-mlx 0.3.0 → 0.4.0

package/ios/Sources/HybridSTT.swift

@@ -0,0 +1,202 @@
+import Foundation
+import NitroModules
+internal import MLX
+internal import MLXAudioSTT
+internal import MLXAudioCore
+
+enum STTError: Error {
+  case notLoaded
+  case notListening
+  case alreadyListening
+}
+
+class HybridSTT: HybridSTTSpec {
+  private var model: GLMASRModel?
+  private var activeTask: Task<String, Error>?
+  private var loadTask: Task<Void, Error>?
+  private var captureManager: AudioCaptureManager?
+
+  var isLoaded: Bool { model != nil }
+  var isTranscribing: Bool { activeTask != nil }
+  var isListening: Bool { captureManager?.isCapturing ?? false }
+  var modelId: String = ""
+
+  private func arrayBufferToMLXArray(_ buffer: ArrayBuffer) -> MLXArray {
+    let count = buffer.size / MemoryLayout<Float>.size
+    let rawPtr = UnsafeRawPointer(buffer.data)
+    let floatPtr = rawPtr.bindMemory(to: Float.self, capacity: count)
+    let floatBuffer = UnsafeBufferPointer(start: floatPtr, count: count)
+    return MLXArray(Array(floatBuffer))
+  }
+
+  func load(modelId: String, options: STTLoadOptions?) throws -> Promise<Void> {
+    self.loadTask?.cancel()
+
+    return Promise.async { [self] in
+      let task = Task { @MainActor in
+        self.activeTask?.cancel()
+        self.activeTask = nil
+        self.model = nil
+        MLX.Memory.clearCache()
+
+        let loadedModel = try await GLMASRModel.fromPretrained(modelId)
+
+        try Task.checkCancellation()
+
+        self.model = loadedModel
+        self.modelId = modelId
+
+        options?.onProgress?(1.0)
+      }
+
+      self.loadTask = task
+      try await task.value
+    }
+  }
+
+  func transcribe(audio: ArrayBuffer) throws -> Promise<String> {
+    guard let model else {
+      throw STTError.notLoaded
+    }
+
+    return Promise.async { [self] in
+      let task = Task<String, Error> {
+        let mlxAudio = self.arrayBufferToMLXArray(audio)
+        let output = model.generate(audio: mlxAudio)
+        return output.text
+      }
+
+      self.activeTask = task
+      defer { self.activeTask = nil }
+
+      return try await task.value
+    }
+  }
+
+  func transcribeStream(
+    audio: ArrayBuffer,
+    onToken: @escaping (_ token: String) -> Void
+  ) throws -> Promise<String> {
+    guard let model else {
+      throw STTError.notLoaded
+    }
+
+    return Promise.async { [self] in
+      let task = Task<String, Error> {
+        let mlxAudio = self.arrayBufferToMLXArray(audio)
+        let stream = model.generateStream(audio: mlxAudio)
+        var finalText = ""
+
+        for try await event in stream {
+          if Task.isCancelled { break }
+
+          switch event {
+          case .token(let token):
+            onToken(token)
+          case .result(let output):
+            finalText = output.text
+          case .info:
+            break
+          }
+        }
+
+        return finalText
+      }
+
+      self.activeTask = task
+      defer { self.activeTask = nil }
+
+      return try await task.value
+    }
+  }
+
+  func startListening() throws -> Promise<Void> {
+    guard model != nil else {
+      throw STTError.notLoaded
+    }
+    guard captureManager == nil || !captureManager!.isCapturing else {
+      throw STTError.alreadyListening
+    }
+
+    return Promise.async { [self] in
+      let manager = AudioCaptureManager()
+      self.captureManager = manager
+      try await manager.startCapturing()
+    }
+  }
+
+  func transcribeBuffer() throws -> Promise<String> {
+    guard let model else {
+      throw STTError.notLoaded
+    }
+    guard let manager = captureManager, manager.isCapturing else {
+      throw STTError.notListening
+    }
+    guard let audio = manager.snapshot() else {
+      return Promise.resolved(withResult: "")
+    }
+
+    return Promise.async { [self] in
+      let task = Task<String, Error> {
+        let output = model.generate(audio: audio)
+        return output.text
+      }
+
+      self.activeTask = task
+      defer { self.activeTask = nil }
+
+      let result = try await task.value
+      MLX.Memory.clearCache()
+      return result
+    }
+  }
+
+  func stopListening() throws -> Promise<String> {
+    guard let model else {
+      throw STTError.notLoaded
+    }
+    guard let manager = captureManager, manager.isCapturing else {
+      throw STTError.notListening
+    }
+
+    let audio = manager.stopCapturing()
+    self.captureManager = nil
+
+    return Promise.async { [self] in
+      let task = Task<String, Error> {
+        let output = model.generate(audio: audio)
+        return output.text
+      }
+
+      self.activeTask = task
+      defer { self.activeTask = nil }
+
+      let result = try await task.value
+      MLX.Memory.clearCache()
+      return result
+    }
+  }
+
+  func stop() throws {
+    activeTask?.cancel()
+    activeTask = nil
+    if let manager = captureManager, manager.isCapturing {
+      _ = manager.stopCapturing()
+    }
+    captureManager = nil
+  }
+
+  func unload() throws {
+    loadTask?.cancel()
+    loadTask = nil
+    activeTask?.cancel()
+    activeTask = nil
+    if let manager = captureManager, manager.isCapturing {
+      _ = manager.stopCapturing()
+    }
+    captureManager = nil
+    model = nil
+    modelId = ""
+    Memory.clearCache()
+  }
+}

package/ios/Sources/HybridTTS.swift

@@ -0,0 +1,145 @@
+import Foundation
+import NitroModules
+internal import MLX
+internal import MLXAudioTTS
+internal import MLXAudioCore
+
+enum TTSError: Error {
+  case notLoaded
+}
+
+class HybridTTS: HybridTTSSpec {
+  private var model: SpeechGenerationModel?
+  private var activeTask: Task<Any, Error>?
+  private var loadTask: Task<Void, Error>?
+
+  var isLoaded: Bool { model != nil }
+  var isGenerating: Bool { activeTask != nil }
+  var modelId: String = ""
+  var sampleRate: Double {
+    Double(model?.sampleRate ?? 24000)
+  }
+
+  private func mlxArrayToArrayBuffer(_ audio: MLXArray) -> ArrayBuffer {
+    let evaluated = audio.asType(.float32)
+    MLX.eval(evaluated)
+    let arrayData = evaluated.asData(access: .copy)
+    let byteSize = arrayData.data.count
+    let buffer = ArrayBuffer.allocate(size: byteSize)
+    arrayData.data.withUnsafeBytes { srcPtr in
+      UnsafeMutableRawPointer(buffer.data).copyMemory(
+        from: srcPtr.baseAddress!,
+        byteCount: byteSize
+      )
+    }
+    return buffer
+  }
+
+  func load(modelId: String, options: TTSLoadOptions?) throws -> Promise<Void> {
+    self.loadTask?.cancel()
+
+    return Promise.async { [self] in
+      let task = Task { @MainActor in
+        self.activeTask?.cancel()
+        self.activeTask = nil
+        self.model = nil
+        MLX.Memory.clearCache()
+
+        let loadedModel = try await TTSModelUtils.loadModel(modelRepo: modelId)
+
+        try Task.checkCancellation()
+
+        self.model = loadedModel
+        self.modelId = modelId
+
+        options?.onProgress?(1.0)
+      }
+
+      self.loadTask = task
+      try await task.value
+    }
+  }
+
+  func generate(
+    text: String,
+    options: TTSGenerateOptions?
+  ) throws -> Promise<ArrayBuffer> {
+    guard let model else {
+      throw TTSError.notLoaded
+    }
+
+    return Promise.async { [self] in
+      let task = Task<Any, Error> {
+        let audio = try await model.generate(
+          text: text,
+          voice: options?.voice,
+          refAudio: nil,
+          refText: nil,
+          language: nil
+        )
+        return self.mlxArrayToArrayBuffer(audio) as Any
+      }
+
+      self.activeTask = task
+      defer { self.activeTask = nil }
+
+      return try await task.value as! ArrayBuffer
+    }
+  }
+
+  func stream(
+    text: String,
+    onAudioChunk: @escaping (ArrayBuffer) -> Void,
+    options: TTSGenerateOptions?
+  ) throws -> Promise<Void> {
+    guard let model else {
+      throw TTSError.notLoaded
+    }
+
+    return Promise.async { [self] in
+      let task = Task<Any, Error> {
+        let stream = model.generateStream(
+          text: text,
+          voice: options?.voice,
+          refAudio: nil,
+          refText: nil,
+          language: nil,
+          generationParameters: model.defaultGenerationParameters
+        )
+
+        for try await event in stream {
+          if Task.isCancelled { break }
+
+          switch event {
+          case .audio(let audio):
+            let buffer = self.mlxArrayToArrayBuffer(audio)
+            onAudioChunk(buffer)
+          case .token, .info:
+            break
+          }
+        }
+        return () as Any
+      }
+
+      self.activeTask = task
+      defer { self.activeTask = nil }
+
+      _ = try await task.value
+    }
+  }
+
+  func stop() throws {
+    activeTask?.cancel()
+    activeTask = nil
+  }
+
+  func unload() throws {
+    loadTask?.cancel()
+    loadTask = nil
+    activeTask?.cancel()
+    activeTask = nil
+    model = nil
+    modelId = ""
+    Memory.clearCache()
+  }
+}

package/ios/Sources/JSONHelpers.swift

@@ -0,0 +1,9 @@
+import Foundation
+
+func dictionaryToJson(_ dict: [String: Any]) -> String {
+    guard let data = try? JSONSerialization.data(withJSONObject: dict),
+          let json = String(data: data, encoding: .utf8) else {
+        return "{}"
+    }
+    return json
+}

package/ios/Sources/ModelDownloader.swift

@@ -12,32 +12,46 @@ actor ModelDownloader: NSObject {
         }
     }
 
+    private let downloadableExtensions: Set<String> = [
+        "json", "safetensors", "txt", "model", "tiktoken", "py"
+    ]
+
+    private func fetchFileList(modelId: String) async throws -> [String] {
+        let urlString = "https://huggingface.co/api/models/\(modelId)"
+        guard let url = URL(string: urlString) else { return [] }
+
+        let (data, _) = try await URLSession.shared.data(from: url)
+        guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any],
+              let siblings = json["siblings"] as? [[String: Any]]
+        else { return [] }
+
+        return siblings.compactMap { $0["rfilename"] as? String }
+            .filter { name in
+                let ext = (name as NSString).pathExtension.lowercased()
+                return downloadableExtensions.contains(ext)
+            }
+    }
+
     func download(
         modelId: String,
         progressCallback: @escaping (Double) -> Void
     ) async throws -> URL {
-        let requiredFiles = [
-            "config.json",
-            "tokenizer.json",
-            "tokenizer_config.json",
-            "model.safetensors"
-        ]
-
+        let files = try await fetchFileList(modelId: modelId)
         let modelDir = getModelDirectory(modelId: modelId)
         try fileManager.createDirectory(at: modelDir, withIntermediateDirectories: true)
 
         log("Model directory: \(modelDir.path)")
-        log("Files to download: \(requiredFiles)")
+        log("Files to download: \(files)")
 
         var downloaded = 0
 
-        for file in requiredFiles {
+        for file in files {
             let destURL = modelDir.appendingPathComponent(file)
 
             if fileManager.fileExists(atPath: destURL.path) {
                 log("File exists, skipping: \(file)")
                 downloaded += 1
-                progressCallback(Double(downloaded) / Double(requiredFiles.count))
+                progressCallback(Double(downloaded) / Double(files.count))
                 continue
             }
 
@@ -69,7 +83,7 @@ actor ModelDownloader: NSObject {
             }
 
             downloaded += 1
-            progressCallback(Double(downloaded) / Double(requiredFiles.count))
+            progressCallback(Double(downloaded) / Double(files.count))
         }
 
         return modelDir
@@ -77,7 +91,7 @@ actor ModelDownloader: NSObject {
 
     func isDownloaded(modelId: String) -> Bool {
         let modelDir = getModelDirectory(modelId: modelId)
-        let requiredFiles = ["config.json", "model.safetensors", "tokenizer.json"]
+        let requiredFiles = ["config.json", "model.safetensors"]
 
         let allExist = requiredFiles.allSatisfy { file in
             fileManager.fileExists(atPath: modelDir.appendingPathComponent(file).path)

package/ios/Sources/StreamEventEmitter.swift

@@ -0,0 +1,132 @@
+import Foundation
+import NitroModules
+
+struct StreamEventEmitter {
+    private let callback: (String) -> Void
+    private let encoder = JSONEncoder()
+
+    init(callback: @escaping (String) -> Void) {
+        self.callback = callback
+    }
+
+    private func emit<T: Encodable>(_ event: T) {
+        guard let data = try? encoder.encode(event),
+              let json = String(data: data, encoding: .utf8) else { return }
+        callback(json)
+    }
+
+    private func timestamp() -> Double {
+        Date().timeIntervalSince1970 * 1000
+    }
+
+    struct GenerationStartEvent: Encodable {
+        let type = "generation_start"
+        let timestamp: Double
+    }
+
+    struct TokenEvent: Encodable {
+        let type = "token"
+        let token: String
+    }
+
+    struct ThinkingStartEvent: Encodable {
+        let type = "thinking_start"
+        let timestamp: Double
+    }
+
+    struct ThinkingChunkEvent: Encodable {
+        let type = "thinking_chunk"
+        let chunk: String
+    }
+
+    struct ThinkingEndEvent: Encodable {
+        let type = "thinking_end"
+        let content: String
+        let timestamp: Double
+    }
+
+    struct ToolCallStartEvent: Encodable {
+        let type = "tool_call_start"
+        let id: String
+        let name: String
+        let arguments: String
+    }
+
+    struct ToolCallExecutingEvent: Encodable {
+        let type = "tool_call_executing"
+        let id: String
+    }
+
+    struct ToolCallCompletedEvent: Encodable {
+        let type = "tool_call_completed"
+        let id: String
+        let result: String
+    }
+
+    struct ToolCallFailedEvent: Encodable {
+        let type = "tool_call_failed"
+        let id: String
+        let error: String
+    }
+
+    struct StatsPayload: Encodable {
+        let tokenCount: Double
+        let tokensPerSecond: Double
+        let timeToFirstToken: Double
+        let totalTime: Double
+    }
+
+    struct GenerationEndEvent: Encodable {
+        let type = "generation_end"
+        let content: String
+        let stats: StatsPayload
+    }
+
+    func emitGenerationStart() {
+        emit(GenerationStartEvent(timestamp: timestamp()))
+    }
+
+    func emitToken(_ token: String) {
+        emit(TokenEvent(token: token))
+    }
+
+    func emitThinkingStart() {
+        emit(ThinkingStartEvent(timestamp: timestamp()))
+    }
+
+    func emitThinkingChunk(_ chunk: String) {
+        emit(ThinkingChunkEvent(chunk: chunk))
+    }
+
+    func emitThinkingEnd(_ content: String) {
+        emit(ThinkingEndEvent(content: content, timestamp: timestamp()))
+    }
+
+    func emitToolCallStart(id: String, name: String, arguments: String) {
+        emit(ToolCallStartEvent(id: id, name: name, arguments: arguments))
+    }
+
+    func emitToolCallExecuting(id: String) {
+        emit(ToolCallExecutingEvent(id: id))
+    }
+
+    func emitToolCallCompleted(id: String, result: String) {
+        emit(ToolCallCompletedEvent(id: id, result: result))
+    }
+
+    func emitToolCallFailed(id: String, error: String) {
+        emit(ToolCallFailedEvent(id: id, error: error))
+    }
+
+    func emitGenerationEnd(content: String, stats: GenerationStats) {
+        emit(GenerationEndEvent(
+            content: content,
+            stats: StatsPayload(
+                tokenCount: stats.tokenCount,
+                tokensPerSecond: stats.tokensPerSecond,
+                timeToFirstToken: stats.timeToFirstToken,
+                totalTime: stats.totalTime
+            )
+        ))
+    }
+}