bleam 0.0.9 → 0.0.11

package/templates/native/ios/GenerationService/BonsaiTextGenerationRunner.swift

@@ -0,0 +1,103 @@
+import Foundation
+import MLX
+
+final class BonsaiTextGenerationRunner: @unchecked Sendable {
+  private let request: BonsaiTextGenerationRequest
+  private let emit: @Sendable (BonsaiTextGenerationEvent) -> Void
+  private let stateLock = NSLock()
+  private var cancellationRequested = false
+  private var generatedTokenIDs: [Int] = []
+
+  init(
+    request: BonsaiTextGenerationRequest,
+    emit: @escaping @Sendable (BonsaiTextGenerationEvent) -> Void
+  ) {
+    self.request = request
+    self.emit = emit
+  }
+
+  func cancel() {
+    stateLock.withLock { cancellationRequested = true }
+  }
+
+  func run(isCancelled externalCancellation: @escaping @Sendable () -> Bool) async {
+    defer { MLX.Memory.clearCache() }
+    do {
+      let profile = try BonsaiTextModelProfile.resolve(request.model)
+      let parameters = try generationParameters(defaults: profile.defaultGenerationParameters)
+      let endTokenIDs = request.endTokenIDs.isEmpty ? profile.defaultEndTokenIDs : Set(request.endTokenIDs)
+      let result = try await BonsaiTextGenerator.generate(
+        profile: profile,
+        modelURL: URL(fileURLWithPath: request.modelPath, isDirectory: true),
+        inputIDs: request.inputIDs,
+        endTokenIDs: endTokenIDs,
+        parameters: parameters,
+        isCancelled: { [weak self] in
+          externalCancellation() || (self?.isCancellationRequested ?? true)
+        },
+        onToken: { [weak self, emit, id = request.id] tokenID in
+          guard let tokenIDs = self?.append(tokenID) else { return false }
+          emit(.init(type: .update, id: id, tokenIDs: tokenIDs))
+          return true
+        }
+      )
+
+      if result.finishReason == .cancelled || shouldCancel(externalCancellation) {
+        emit(.init(type: .canceled, id: request.id, tokenIDs: result.tokenIDs, finishReason: GenerationFinishReason.cancelled.rawValue))
+      } else {
+        emit(.init(type: .completed, id: request.id, tokenIDs: result.tokenIDs, finishReason: result.finishReason.rawValue))
+      }
+    } catch is CancellationError {
+      emit(.init(type: .canceled, id: request.id, tokenIDs: generatedTokens, finishReason: GenerationFinishReason.cancelled.rawValue))
+    } catch {
+      emit(.init(type: .failed, id: request.id, tokenIDs: generatedTokens, message: error.localizedDescription))
+    }
+  }
+
+  private var isCancellationRequested: Bool {
+    stateLock.withLock { cancellationRequested }
+  }
+
+  private var generatedTokens: [Int] {
+    stateLock.withLock { generatedTokenIDs }
+  }
+
+  private func append(_ tokenID: Int) -> [Int] {
+    stateLock.withLock {
+      generatedTokenIDs.append(tokenID)
+      return generatedTokenIDs
+    }
+  }
+
+  private func shouldCancel(_ externalCancellation: @escaping @Sendable () -> Bool) -> Bool {
+    Task.isCancelled || externalCancellation() || isCancellationRequested
+  }
+
+  private func generationParameters(defaults: GenerationParameters) throws -> GenerationParameters {
+    let seed: UInt64?
+    if let requestedSeed = request.seed {
+      guard let converted = UInt64(exactly: requestedSeed) else {
+        throw BonsaiTextError.invalidSeed(requestedSeed)
+      }
+      seed = converted
+    } else {
+      seed = defaults.seed
+    }
+
+    let parameters = GenerationParameters(
+      maximumNewTokens: request.maximumNewTokens ?? defaults.maximumNewTokens,
+      temperature: request.temperature.map(Float.init) ?? defaults.temperature,
+      topP: request.topP.map(Float.init) ?? defaults.topP,
+      topK: request.topK ?? defaults.topK,
+      repetitionPenalty: request.repetitionPenalty.map(Float.init) ?? defaults.repetitionPenalty,
+      repetitionContextSize: request.repetitionContextSize ?? defaults.repetitionContextSize,
+      seed: seed
+    )
+    try TokenSampler.validate(parameters)
+    guard parameters.temperature.isFinite, parameters.topP.isFinite,
+          parameters.repetitionPenalty.isFinite else {
+      throw BonsaiTextError.invalidGenerationParameters("floating-point values must be finite")
+    }
+    return parameters
+  }
+}

package/templates/native/ios/GenerationService/ImageGenerationRunner.swift

@@ -7,6 +7,7 @@ import UniformTypeIdentifiers
 final class ImageGenerationRunner: @unchecked Sendable {
   private let request: GenerationWorkerRequest
   private let emit: @Sendable (GenerationWorkerEvent) -> Void
+  private let stateLock = NSLock()
   private var isCancelled = false
 
   init(
@@ -18,7 +19,7 @@ final class ImageGenerationRunner: @unchecked Sendable {
   }
 
   func cancel() {
-    isCancelled = true
+    stateLock.withLock { isCancelled = true }
   }
 
   func run(isCancelled externalCancellation: @escaping @Sendable () -> Bool) async throws {
@@ -88,10 +89,13 @@ final class ImageGenerationRunner: @unchecked Sendable {
       // First service version is deliberately memory-first: the runner releases
       // text encoder, transformer, and VAE stages instead of keeping a warm
       // pipeline. Add warm caching later only behind explicit idle trimming.
-      let runner = FluxImageRunner(paths: FluxImageRunnerPaths(
-        transformer: URL(fileURLWithPath: request.transformerPath, isDirectory: true),
-        vae: URL(fileURLWithPath: request.vaePath, isDirectory: true)
-      ))
+       let runner = FluxImageRunner(
+         paths: FluxImageRunnerPaths(
+           transformer: URL(fileURLWithPath: request.transformerPath, isDirectory: true),
+           vae: URL(fileURLWithPath: request.vaePath, isDirectory: true)
+         ),
+         transformerProfile: transformerProfile()
+       )
       var fluxStage = "denoise"
       do {
         return try runner.generate(
@@ -147,11 +151,13 @@ final class ImageGenerationRunner: @unchecked Sendable {
   }
 
   private func shouldCancel(_ externalCancellation: @escaping @Sendable () -> Bool) -> Bool {
-    Task.isCancelled || isCancelled || externalCancellation()
+    Task.isCancelled || stateLock.withLock { isCancelled } || externalCancellation()
   }
 
   private func validateRequest() throws {
-    guard request.model == "flux2-klein-4b-8bit-abliterated" else {
+    guard request.model == "flux2-klein-4b-8bit-abliterated"
+      || request.model == "ternary-bonsai-image-4b"
+    else {
       throw GenerationWorkerFailure(message: "Unsupported image model: \(request.model)")
     }
     guard !request.prompt.isEmpty else {
@@ -264,6 +270,12 @@ final class ImageGenerationRunner: @unchecked Sendable {
     }
   }
 
+  private func transformerProfile() -> Flux2TransformerRepoProfile {
+    request.model == "ternary-bonsai-image-4b"
+      ? ModelRegistry.TransformerRepo.prismMLBonsaiKlein4B2bit.profile
+      : ModelRegistry.TransformerRepo.moxinOrgKlein4B8bit.profile
+  }
+
   private static func validateSafetensors(atPath path: String) throws {
     let url = URL(fileURLWithPath: path)
     let fileSize: UInt64

package/templates/native/ios/GenerationService/main.swift

@@ -2,68 +2,130 @@ import Foundation
 import MLX
 
 private final class GenerationRuntime: @unchecked Sendable {
+  private enum OperationKind {
+    case image
+    case text
+  }
+
   private let modelWarmTimeout: Duration = .seconds(60)
+  private let stateLock = NSLock()
+  private var activeID: String?
+  private var activeKind: OperationKind?
   private var activeTask: Task<Void, Never>?
-  private var activeRunner: ImageGenerationRunner?
+  private var activeImageRunner: ImageGenerationRunner?
+  private var activeTextRunner: BonsaiTextGenerationRunner?
   private var trimTask: Task<Void, Never>?
   private var isCancelled = false
+  private var operationEpoch: UInt64 = 0
+  private var isShuttingDown = false
 
   var isRunning: Bool {
-    activeTask != nil
+    stateLock.withLock { activeTask != nil }
   }
 
-  func start(
+  func startImage(
     request: GenerationWorkerRequest,
     emit: @escaping @Sendable (GenerationWorkerEvent) -> Void
   ) -> Bool {
-    guard activeTask == nil else { return false }
+    stateLock.withLock {
+      guard activeTask == nil, !isShuttingDown else { return false }
+      trimTask?.cancel()
+      trimTask = nil
+      isCancelled = false
+      operationEpoch &+= 1
+      activeID = request.jobID.uuidString
+      activeKind = .image
+      let runner = ImageGenerationRunner(request: request, emit: emit)
+      activeImageRunner = runner
+      activeTask = Task { [weak self] in
+        await self?.runImage(request, runner: runner, emit: emit)
+      }
+      return true
+    }
+  }
+
+  func startText(
+    request: BonsaiTextGenerationRequest,
+    emit: @escaping @Sendable (BonsaiTextGenerationEvent) -> Void
+  ) -> Bool {
+    stateLock.withLock {
+      guard activeTask == nil, !isShuttingDown else { return false }
+      trimTask?.cancel()
+      trimTask = nil
+      isCancelled = false
+      operationEpoch &+= 1
+      activeID = request.id
+      activeKind = .text
+      let runner = BonsaiTextGenerationRunner(request: request, emit: emit)
+      activeTextRunner = runner
+      activeTask = Task { [weak self] in
+        await self?.runText(request, runner: runner)
+      }
+      return true
+    }
+  }
+
+  func cancelImage() -> Bool {
+    let operation = stateLock.withLock { () -> (Task<Void, Never>?, ImageGenerationRunner?)? in
+      guard activeKind == .image else { return nil }
+      isCancelled = true
+      trimTask?.cancel()
+      trimTask = nil
+      return (activeTask, activeImageRunner)
+    }
+    guard let operation else { return false }
+    operation.0?.cancel()
+    operation.1?.cancel()
+    return true
+  }
 
-    trimTask?.cancel()
-    trimTask = nil
-    isCancelled = false
-    activeTask = Task { [weak self] in
-      await self?.run(request, emit: emit)
+  func cancelText(id: String) -> Bool {
+    let operation = stateLock.withLock { () -> (Task<Void, Never>?, BonsaiTextGenerationRunner?)? in
+      guard activeKind == .text, activeID == id else { return nil }
+      isCancelled = true
+      trimTask?.cancel()
+      trimTask = nil
+      return (activeTask, activeTextRunner)
     }
+    guard let operation else { return false }
+    operation.0?.cancel()
+    operation.1?.cancel()
     return true
   }
 
-  func cancel() {
-    isCancelled = true
-    trimTask?.cancel()
-    trimTask = nil
-    activeTask?.cancel()
-    activeRunner?.cancel()
+  func beginShutdown() -> Bool {
+    stateLock.withLock {
+      guard activeTask == nil, !isShuttingDown else { return false }
+      isShuttingDown = true
+      operationEpoch &+= 1
+      trimTask?.cancel()
+      trimTask = nil
+      return true
+    }
   }
 
   func trimModelsNow() async {
-    trimTask?.cancel()
-    trimTask = nil
-    // We currently reload every stage per generation, but keep the explicit
-    // trim hook from ImageStudio's service architecture. Warm model caching
-    // should be added here later, never back in the app process.
-    MLX.Memory.clearCache()
-    await Task.yield()
-    MLX.Memory.clearCache()
+    stateLock.withLock {
+      trimTask?.cancel()
+      trimTask = nil
+      MLX.Memory.clearCache()
+      MLX.Memory.clearCache()
+    }
     print("[GenerationService] trim active=\(MLX.Memory.activeMemory / 1_048_576)MB cache=\(MLX.Memory.cacheMemory / 1_048_576)MB")
   }
 
-  private func run(
+  private func runImage(
     _ request: GenerationWorkerRequest,
+    runner: ImageGenerationRunner,
     emit: @escaping @Sendable (GenerationWorkerEvent) -> Void
   ) async {
-    let runner = ImageGenerationRunner(request: request, emit: emit)
-    activeRunner = runner
     defer {
-      if activeRunner === runner {
-        activeRunner = nil
-      }
-      activeTask = nil
-      scheduleTrimAfterIdle()
+      finish(id: request.jobID.uuidString, imageRunner: runner)
     }
 
     do {
       try await runner.run(isCancelled: { [weak self] in
-        Task.isCancelled || (self?.isCancelled ?? false)
+        Task.isCancelled || (self?.cancellationRequested ?? false)
       })
     } catch is CancellationError {
       emit(.init(type: .canceled))
@@ -76,17 +138,59 @@ private final class GenerationRuntime: @unchecked Sendable {
     }
   }
 
+  private func runText(
+    _ request: BonsaiTextGenerationRequest,
+    runner: BonsaiTextGenerationRunner
+  ) async {
+    defer { finish(id: request.id, textRunner: runner) }
+    await runner.run(isCancelled: { [weak self] in
+      Task.isCancelled || (self?.cancellationRequested ?? true)
+    })
+  }
+
+  private var cancellationRequested: Bool {
+    stateLock.withLock { isCancelled }
+  }
+
+  private func finish(
+    id: String,
+    imageRunner: ImageGenerationRunner? = nil,
+    textRunner: BonsaiTextGenerationRunner? = nil
+  ) {
+    let finished = stateLock.withLock {
+      guard activeID == id else { return false }
+      if let imageRunner, activeImageRunner !== imageRunner { return false }
+      if let textRunner, activeTextRunner !== textRunner { return false }
+      activeID = nil
+      activeKind = nil
+      activeImageRunner = nil
+      activeTextRunner = nil
+      activeTask = nil
+      isCancelled = false
+      return true
+    }
+    if finished { scheduleTrimAfterIdle() }
+  }
+
   private func scheduleTrimAfterIdle() {
-    trimTask?.cancel()
-    trimTask = Task { [weak self] in
-      guard let self else { return }
-      do {
-        try await Task.sleep(for: modelWarmTimeout)
-      } catch {
-        return
+    stateLock.withLock {
+      trimTask?.cancel()
+      let scheduledEpoch = operationEpoch
+      trimTask = Task { [weak self] in
+        guard let self else { return }
+        do {
+          try await Task.sleep(for: modelWarmTimeout)
+        } catch {
+          return
+        }
+        guard !Task.isCancelled else { return }
+        self.stateLock.withLock {
+          guard self.activeTask == nil, self.operationEpoch == scheduledEpoch else { return }
+          self.trimTask = nil
+          MLX.Memory.clearCache()
+          MLX.Memory.clearCache()
+        }
       }
-      guard !Task.isCancelled else { return }
-      await self.trimModelsNow()
     }
     print("[GenerationService] model trim scheduled 60s")
   }
@@ -94,10 +198,11 @@ private final class GenerationRuntime: @unchecked Sendable {
 
 private final class GenerationService: NSObject, GenerationServiceProtocol {
   private weak var connection: NSXPCConnection?
-  private let runtime = GenerationRuntime()
+  private let runtime: GenerationRuntime
 
-  init(connection: NSXPCConnection) {
+  init(connection: NSXPCConnection, runtime: GenerationRuntime) {
     self.connection = connection
+    self.runtime = runtime
   }
 
   func startGeneration(requestData: Data, reply: @escaping (Bool) -> Void) {
@@ -107,14 +212,27 @@ private final class GenerationService: NSObject, GenerationServiceProtocol {
       return
     }
 
-    reply(runtime.start(request: request) { [weak self] event in
+    reply(runtime.startImage(request: request) { [weak self] event in
       self?.emit(event)
     })
   }
 
   func cancelGeneration(reply: @escaping (Bool) -> Void) {
-    runtime.cancel()
-    reply(true)
+    reply(runtime.cancelImage())
+  }
+
+  func startTextGeneration(requestData: Data, reply: @escaping (Bool) -> Void) {
+    guard let request = try? JSONDecoder().decode(BonsaiTextGenerationRequest.self, from: requestData) else {
+      reply(false)
+      return
+    }
+    reply(runtime.startText(request: request) { [weak self] event in
+      self?.emitText(event)
+    })
+  }
+
+  func cancelTextGeneration(id: String, reply: @escaping (Bool) -> Void) {
+    reply(runtime.cancelText(id: id))
   }
 
   func memoryTelemetry(reply: @escaping (Data?) -> Void) {
@@ -126,7 +244,7 @@ private final class GenerationService: NSObject, GenerationServiceProtocol {
   }
 
   func shutdown(reply: @escaping (Bool) -> Void) {
-    guard !runtime.isRunning else {
+    guard runtime.beginShutdown() else {
       reply(false)
       return
     }
@@ -142,13 +260,21 @@ private final class GenerationService: NSObject, GenerationServiceProtocol {
     let client = connection?.remoteObjectProxy as? GenerationServiceClientProtocol
     client?.generationWorkerDidEmit(eventData: data)
   }
+
+  private func emitText(_ event: BonsaiTextGenerationEvent) {
+    guard let data = try? JSONEncoder().encode(event) else { return }
+    let client = connection?.remoteObjectProxy as? GenerationServiceClientProtocol
+    client?.textGenerationWorkerDidEmit(eventData: data)
+  }
 }
 
 private final class GenerationServiceDelegate: NSObject, NSXPCListenerDelegate {
+  private let runtime = GenerationRuntime()
+
   func listener(_ listener: NSXPCListener, shouldAcceptNewConnection connection: NSXPCConnection) -> Bool {
     connection.exportedInterface = NSXPCInterface(with: GenerationServiceProtocol.self)
     connection.remoteObjectInterface = NSXPCInterface(with: GenerationServiceClientProtocol.self)
-    connection.exportedObject = GenerationService(connection: connection)
+    connection.exportedObject = GenerationService(connection: connection, runtime: runtime)
     connection.resume()
     return true
   }

package/templates/native/ios/Shared/Generation/GenerationServiceProtocol.swift

@@ -3,10 +3,13 @@ import Foundation
 @objc protocol GenerationServiceProtocol {
   func startGeneration(requestData: Data, reply: @escaping (Bool) -> Void)
   func cancelGeneration(reply: @escaping (Bool) -> Void)
+  func startTextGeneration(requestData: Data, reply: @escaping (Bool) -> Void)
+  func cancelTextGeneration(id: String, reply: @escaping (Bool) -> Void)
   func memoryTelemetry(reply: @escaping (Data?) -> Void)
   func shutdown(reply: @escaping (Bool) -> Void)
 }
 
 @objc protocol GenerationServiceClientProtocol {
   func generationWorkerDidEmit(eventData: Data)
+  func textGenerationWorkerDidEmit(eventData: Data)
 }

package/templates/native/ios/Shared/Generation/GenerationWorkerProtocol.swift

@@ -122,3 +122,47 @@ struct GenerationWorkerLoRA: Codable, Sendable {
   let path: String
   let scale: Double
 }
+
+struct BonsaiTextGenerationRequest: Codable, Sendable {
+  let id: String
+  let model: String
+  let modelPath: String
+  let inputIDs: [Int]
+  let endTokenIDs: [Int]
+  let maximumNewTokens: Int?
+  let temperature: Double?
+  let topP: Double?
+  let topK: Int?
+  let repetitionPenalty: Double?
+  let repetitionContextSize: Int?
+  let seed: Int?
+}
+
+struct BonsaiTextGenerationEvent: Codable, Sendable {
+  enum EventType: String, Codable, Sendable {
+    case update
+    case completed
+    case canceled
+    case failed
+  }
+
+  let type: EventType
+  let id: String
+  let tokenIDs: [Int]
+  let finishReason: String?
+  let message: String?
+
+  init(
+    type: EventType,
+    id: String,
+    tokenIDs: [Int] = [],
+    finishReason: String? = nil,
+    message: String? = nil
+  ) {
+    self.type = type
+    self.id = id
+    self.tokenIDs = tokenIDs
+    self.finishReason = finishReason
+    self.message = message
+  }
+}