bleam 0.0.8 → 0.0.9

package/templates/native/ios/Bleam/AI/Flux2/Transformer/Flux2TransformerBlock.swift

@@ -0,0 +1,362 @@
+// Flux2TransformerBlock.swift - Double-Stream Transformer Block
+// Copyright 2025 Vincent Gourbin
+
+import Foundation
+import MLX
+import MLXNN
+
+/// Double-Stream Transformer Block for Flux.2
+///
+/// Processes image and text hidden states in parallel streams with:
+/// 1. Separate LayerNorm for each modality
+/// 2. Joint attention across both modalities
+/// 3. Separate FeedForward for each modality
+/// 4. AdaLN modulation from timestep embeddings
+///
+/// There are 8 such blocks in Flux.2.
+public class Flux2TransformerBlock: Module, @unchecked Sendable {
+    let dim: Int
+    let numHeads: Int
+    let headDim: Int
+
+    // Layer norms for image stream
+    let norm1: LayerNorm
+    let norm2: LayerNorm
+
+    // Layer norms for text stream
+    let norm1Context: LayerNorm
+    let norm2Context: LayerNorm
+
+    // Joint attention
+    let attn: Flux2Attention
+
+    // FeedForward networks
+    let ff: Flux2FeedForward
+    let ffContext: Flux2FeedForward
+
+    /// Initialize Double-Stream Block
+    /// - Parameters:
+    ///   - dim: Model dimension (6144)
+    ///   - numHeads: Number of attention heads (48)
+    ///   - headDim: Dimension per head (128)
+    ///   - mlpRatio: MLP expansion ratio
+    public init(
+        dim: Int,
+        numHeads: Int,
+        headDim: Int,
+        mlpRatio: Float = 3.0
+    ) {
+        self.dim = dim
+        self.numHeads = numHeads
+        self.headDim = headDim
+
+        let mlpHidden = Int(Float(dim) * mlpRatio)
+
+        // Image stream norms
+        self.norm1 = LayerNorm(dimensions: dim, eps: 1e-6, affine: false)
+        self.norm2 = LayerNorm(dimensions: dim, eps: 1e-6, affine: false)
+
+        // Text stream norms
+        self.norm1Context = LayerNorm(dimensions: dim, eps: 1e-6, affine: false)
+        self.norm2Context = LayerNorm(dimensions: dim, eps: 1e-6, affine: false)
+
+        // Joint attention
+        self.attn = Flux2Attention(dim: dim, numHeads: numHeads, headDim: headDim)
+
+        // FeedForward
+        self.ff = Flux2FeedForward(dim: dim, innerDim: mlpHidden)
+        self.ffContext = Flux2FeedForward(dim: dim, innerDim: mlpHidden)
+    }
+
+    /// Forward pass
+    /// - Parameters:
+    ///   - hiddenStates: Image hidden states [B, S_img, dim]
+    ///   - encoderHiddenStates: Text hidden states [B, S_txt, dim]
+    ///   - temb: Timestep embedding [B, dim]
+    ///   - rotaryEmb: Optional RoPE embeddings
+    ///   - imgModParams: Modulation params for image stream (2 sets: attn, ffn)
+    ///   - txtModParams: Modulation params for text stream (2 sets: attn, ffn)
+    /// - Returns: Updated (text hidden states, image hidden states)
+    public func callAsFunction(
+        hiddenStates: MLXArray,
+        encoderHiddenStates: MLXArray,
+        temb: MLXArray,
+        rotaryEmb: (cos: MLXArray, sin: MLXArray)? = nil,
+        imgModParams: [ModulationParams]? = nil,
+        txtModParams: [ModulationParams]? = nil
+    ) -> (encoderHiddenStates: MLXArray, hiddenStates: MLXArray) {
+        // Store residuals
+        let residualImg = hiddenStates
+        let residualTxt = encoderHiddenStates
+
+        // --- Attention Block ---
+
+        // Normalize image hidden states
+        var imgNorm = norm1(hiddenStates)
+
+        // Normalize text hidden states
+        var txtNorm = norm1Context(encoderHiddenStates)
+
+        // Apply modulation if provided
+        if let imgMod = imgModParams, imgMod.count >= 1 {
+            imgNorm = applyModulation(imgNorm, shift: imgMod[0].shift, scale: imgMod[0].scale)
+        }
+        if let txtMod = txtModParams, txtMod.count >= 1 {
+            txtNorm = applyModulation(txtNorm, shift: txtMod[0].shift, scale: txtMod[0].scale)
+        }
+
+        // Joint attention
+        let (imgAttnOut, txtAttnOut) = attn(
+            hiddenStates: imgNorm,
+            encoderHiddenStates: txtNorm,
+            rotaryEmb: rotaryEmb
+        )
+
+        // Apply gate and add residual
+        var imgOut: MLXArray
+        var txtOut: MLXArray
+
+        if let imgMod = imgModParams, imgMod.count >= 1 {
+            imgOut = residualImg + applyGate(imgAttnOut, gate: imgMod[0].gate)
+        } else {
+            imgOut = residualImg + imgAttnOut
+        }
+
+        if let txtMod = txtModParams, txtMod.count >= 1 {
+            txtOut = residualTxt + applyGate(txtAttnOut, gate: txtMod[0].gate)
+        } else {
+            txtOut = residualTxt + txtAttnOut
+        }
+
+        // --- FeedForward Block ---
+
+        // Store new residuals
+        let residualImg2 = imgOut
+        let residualTxt2 = txtOut
+
+        // Normalize
+        var imgNorm2 = norm2(imgOut)
+        var txtNorm2 = norm2Context(txtOut)
+
+        // Apply modulation if provided
+        if let imgMod = imgModParams, imgMod.count >= 2 {
+            imgNorm2 = applyModulation(imgNorm2, shift: imgMod[1].shift, scale: imgMod[1].scale)
+        }
+        if let txtMod = txtModParams, txtMod.count >= 2 {
+            txtNorm2 = applyModulation(txtNorm2, shift: txtMod[1].shift, scale: txtMod[1].scale)
+        }
+
+        // FeedForward
+        let imgFFOut = ff(imgNorm2)
+        let txtFFOut = ffContext(txtNorm2)
+
+        // Apply gate and add residual
+        if let imgMod = imgModParams, imgMod.count >= 2 {
+            imgOut = residualImg2 + applyGate(imgFFOut, gate: imgMod[1].gate)
+        } else {
+            imgOut = residualImg2 + imgFFOut
+        }
+
+        if let txtMod = txtModParams, txtMod.count >= 2 {
+            txtOut = residualTxt2 + applyGate(txtFFOut, gate: txtMod[1].gate)
+        } else {
+            txtOut = residualTxt2 + txtFFOut
+        }
+
+        // Return order matches diffusers: (encoder_hidden_states, hidden_states)
+        return (encoderHiddenStates: txtOut, hiddenStates: imgOut)
+    }
+
+    /// Forward pass with cached text K/V (for KV cache acceleration)
+    /// - Parameters:
+    ///   - hiddenStates: Image hidden states [B, S_img, dim]
+    ///   - encoderHiddenStates: Text hidden states [B, S_txt, dim] (used only for Q projection)
+    ///   - temb: Timestep embedding [B, dim]
+    ///   - cachedTextKV: Pre-computed text K and V
+    ///   - rotaryEmb: Optional RoPE embeddings
+    ///   - imgModParams: Modulation params for image stream
+    ///   - txtModParams: Modulation params for text stream
+    /// - Returns: Updated (text hidden states, image hidden states)
+    public func callAsFunction(
+        hiddenStates: MLXArray,
+        encoderHiddenStates: MLXArray,
+        temb: MLXArray,
+        cachedTextKV: (k: MLXArray, v: MLXArray),
+        rotaryEmb: (cos: MLXArray, sin: MLXArray)? = nil,
+        imgModParams: [ModulationParams]? = nil,
+        txtModParams: [ModulationParams]? = nil
+    ) -> (encoderHiddenStates: MLXArray, hiddenStates: MLXArray) {
+        let residualImg = hiddenStates
+        let residualTxt = encoderHiddenStates
+
+        // --- Attention Block ---
+
+        var imgNorm = norm1(hiddenStates)
+        var txtNorm = norm1Context(encoderHiddenStates)
+
+        if let imgMod = imgModParams, imgMod.count >= 1 {
+            imgNorm = applyModulation(imgNorm, shift: imgMod[0].shift, scale: imgMod[0].scale)
+        }
+        if let txtMod = txtModParams, txtMod.count >= 1 {
+            txtNorm = applyModulation(txtNorm, shift: txtMod[0].shift, scale: txtMod[0].scale)
+        }
+
+        // Use cached text K/V
+        let (imgAttnOut, txtAttnOut) = attn(
+            hiddenStates: imgNorm,
+            encoderHiddenStates: txtNorm,
+            cachedTextKV: cachedTextKV,
+            rotaryEmb: rotaryEmb
+        )
+
+        var imgOut: MLXArray
+        var txtOut: MLXArray
+
+        if let imgMod = imgModParams, imgMod.count >= 1 {
+            imgOut = residualImg + applyGate(imgAttnOut, gate: imgMod[0].gate)
+        } else {
+            imgOut = residualImg + imgAttnOut
+        }
+
+        if let txtMod = txtModParams, txtMod.count >= 1 {
+            txtOut = residualTxt + applyGate(txtAttnOut, gate: txtMod[0].gate)
+        } else {
+            txtOut = residualTxt + txtAttnOut
+        }
+
+        // --- FeedForward Block ---
+
+        let residualImg2 = imgOut
+        let residualTxt2 = txtOut
+
+        var imgNorm2 = norm2(imgOut)
+        var txtNorm2 = norm2Context(txtOut)
+
+        if let imgMod = imgModParams, imgMod.count >= 2 {
+            imgNorm2 = applyModulation(imgNorm2, shift: imgMod[1].shift, scale: imgMod[1].scale)
+        }
+        if let txtMod = txtModParams, txtMod.count >= 2 {
+            txtNorm2 = applyModulation(txtNorm2, shift: txtMod[1].shift, scale: txtMod[1].scale)
+        }
+
+        let imgFFOut = ff(imgNorm2)
+        let txtFFOut = ffContext(txtNorm2)
+
+        if let imgMod = imgModParams, imgMod.count >= 2 {
+            imgOut = residualImg2 + applyGate(imgFFOut, gate: imgMod[1].gate)
+        } else {
+            imgOut = residualImg2 + imgFFOut
+        }
+
+        if let txtMod = txtModParams, txtMod.count >= 2 {
+            txtOut = residualTxt2 + applyGate(txtFFOut, gate: txtMod[1].gate)
+        } else {
+            txtOut = residualTxt2 + txtFFOut
+        }
+
+        return (encoderHiddenStates: txtOut, hiddenStates: imgOut)
+    }
+
+    /// Forward pass for text tokens only (for KV cache building)
+    /// - Parameters:
+    ///   - encoderHiddenStates: Text hidden states [B, S_txt, dim]
+    ///   - temb: Timestep embedding [B, dim]
+    ///   - txtModParams: Modulation params for text stream
+    /// - Returns: Updated text hidden states
+    public func forwardTextOnly(
+        encoderHiddenStates: MLXArray,
+        temb: MLXArray,
+        txtModParams: [ModulationParams]? = nil
+    ) -> MLXArray {
+        let residualTxt = encoderHiddenStates
+
+        // Normalize
+        var txtNorm = norm1Context(encoderHiddenStates)
+        if let txtMod = txtModParams, txtMod.count >= 1 {
+            txtNorm = applyModulation(txtNorm, shift: txtMod[0].shift, scale: txtMod[0].scale)
+        }
+
+        // Text self-attention (using only text K/V, no image interaction)
+        // For KV cache building, we approximate by doing a self-attention on text only
+        let txtAttnOut = attn.forwardTextSelfAttention(
+            encoderHiddenStates: txtNorm
+        )
+
+        var txtOut: MLXArray
+        if let txtMod = txtModParams, txtMod.count >= 1 {
+            txtOut = residualTxt + applyGate(txtAttnOut, gate: txtMod[0].gate)
+        } else {
+            txtOut = residualTxt + txtAttnOut
+        }
+
+        // FeedForward
+        let residualTxt2 = txtOut
+        var txtNorm2 = norm2Context(txtOut)
+        if let txtMod = txtModParams, txtMod.count >= 2 {
+            txtNorm2 = applyModulation(txtNorm2, shift: txtMod[1].shift, scale: txtMod[1].scale)
+        }
+
+        let txtFFOut = ffContext(txtNorm2)
+        if let txtMod = txtModParams, txtMod.count >= 2 {
+            txtOut = residualTxt2 + applyGate(txtFFOut, gate: txtMod[1].gate)
+        } else {
+            txtOut = residualTxt2 + txtFFOut
+        }
+
+        return txtOut
+    }
+}
+
+/// Stack of Double-Stream Transformer Blocks
+public class Flux2TransformerBlocks: Module, @unchecked Sendable {
+    let blocks: [Flux2TransformerBlock]
+
+    public init(
+        numBlocks: Int,
+        dim: Int,
+        numHeads: Int,
+        headDim: Int,
+        mlpRatio: Float = 3.0
+    ) {
+        self.blocks = (0..<numBlocks).map { _ in
+            Flux2TransformerBlock(
+                dim: dim,
+                numHeads: numHeads,
+                headDim: headDim,
+                mlpRatio: mlpRatio
+            )
+        }
+    }
+
+    public func callAsFunction(
+        hiddenStates: MLXArray,
+        encoderHiddenStates: MLXArray,
+        temb: MLXArray,
+        rotaryEmb: (cos: MLXArray, sin: MLXArray)? = nil,
+        imgModulation: Flux2Modulation? = nil,
+        txtModulation: Flux2Modulation? = nil
+    ) -> (encoderHiddenStates: MLXArray, hiddenStates: MLXArray) {
+        var imgHS = hiddenStates
+        var txtHS = encoderHiddenStates
+
+        for block in blocks {
+            // Get modulation params if modulators provided
+            let imgParams = imgModulation?(temb)
+            let txtParams = txtModulation?(temb)
+
+            let (newTxt, newImg) = block(
+                hiddenStates: imgHS,
+                encoderHiddenStates: txtHS,
+                temb: temb,
+                rotaryEmb: rotaryEmb,
+                imgModParams: imgParams,
+                txtModParams: txtParams
+            )
+
+            imgHS = newImg
+            txtHS = newTxt
+        }
+
+        return (encoderHiddenStates: txtHS, hiddenStates: imgHS)
+    }
+}

package/templates/native/ios/Bleam/AI/Flux2/Utils/Flux2Debug.swift

@@ -0,0 +1,101 @@
+// Flux2Debug.swift - Debug and logging utilities
+// Copyright 2025 Vincent Gourbin
+
+import Foundation
+
+/// Debug utilities for Flux.2
+public enum Flux2Debug {
+    public typealias LogSink = @Sendable (Level, String) -> Void
+
+    /// Enable/disable debug logging
+    nonisolated(unsafe) public static var enabled: Bool = true
+
+    nonisolated(unsafe) private static var logSink: LogSink?
+
+    /// Log level
+    public enum Level: Int, Comparable, Sendable {
+        case verbose = 0  // Detailed debug info
+        case info = 1     // General progress
+        case warning = 2  // Warnings
+        case error = 3    // Errors only
+
+        public static func < (lhs: Level, rhs: Level) -> Bool {
+            lhs.rawValue < rhs.rawValue
+        }
+    }
+
+    /// Minimum log level to display (default: info - show progress and above)
+    nonisolated(unsafe) public static var minLevel: Level = .info
+
+    /// Enable debug mode (shows all logs including verbose)
+    public static func enableDebugMode() {
+        enabled = true
+        minLevel = .verbose
+    }
+
+    /// Set to normal mode (only show warnings and errors)
+    public static func setNormalMode() {
+        enabled = true
+        minLevel = .warning
+    }
+
+    public static func setLogSink(_ sink: LogSink?) {
+        logSink = sink
+    }
+
+    /// Log a debug message
+    public static func log(_ message: String, level: Level = .info) {
+        guard enabled, level >= minLevel else { return }
+
+        let prefix: String
+        switch level {
+        case .verbose:
+            prefix = "[Flux2:V]"
+        case .info:
+            prefix = "[Flux2]"
+        case .warning:
+            prefix = "[Flux2:W]"
+        case .error:
+            prefix = "[Flux2:E]"
+        }
+
+        print("\(prefix) \(message)")
+        fflush(stdout)
+        logSink?(level, message)
+    }
+
+    /// Convenience methods
+    public static func verbose(_ message: String) {
+        log(message, level: .verbose)
+    }
+
+    public static func info(_ message: String) {
+        log(message, level: .info)
+    }
+
+    public static func warning(_ message: String) {
+        log(message, level: .warning)
+    }
+
+    public static func error(_ message: String) {
+        log(message, level: .error)
+    }
+
+    /// Log with timing
+    public static func timed<T>(_ label: String, _ block: () throws -> T) rethrows -> T {
+        let start = CFAbsoluteTimeGetCurrent()
+        let result = try block()
+        let elapsed = CFAbsoluteTimeGetCurrent() - start
+        log("\(label): \(String(format: "%.3f", elapsed))s")
+        return result
+    }
+
+    /// Async version of timed
+    public static func timedAsync<T>(_ label: String, _ block: () async throws -> T) async rethrows -> T {
+        let start = CFAbsoluteTimeGetCurrent()
+        let result = try await block()
+        let elapsed = CFAbsoluteTimeGetCurrent() - start
+        log("\(label): \(String(format: "%.3f", elapsed))s")
+        return result
+    }
+}

package/templates/native/ios/Bleam/AI/Flux2/Utils/MLXCheckpoint.swift

@@ -0,0 +1,118 @@
+// MLXCheckpoint.swift - Wrapper for mlx_checkpoint C function
+// Enables gradient checkpointing to reduce memory usage during training
+
+import Cmlx
+import MLX
+
+/// Wraps a function with gradient checkpointing.
+///
+/// During the forward pass, intermediate activations are not stored.
+/// During the backward pass, they are recomputed as needed.
+/// This trades compute time for memory savings (~30-50% less memory).
+///
+/// Usage:
+/// ```swift
+/// let checkpointedBlock = checkpoint { inputs in
+///     return myTransformerBlock(inputs)
+/// }
+/// let output = checkpointedBlock(inputs)
+/// ```
+///
+/// - Parameter f: The function to wrap with checkpointing
+/// - Returns: A checkpointed version of the function
+public func checkpoint(_ f: @escaping ([MLXArray]) -> [MLXArray]) -> ([MLXArray]) -> [MLXArray] {
+    // Create C closure from Swift function
+    let inputClosure = new_mlx_closure(f)
+
+    // Create checkpointed closure
+    var checkpointedClosure = mlx_closure_new()
+    let status = mlx_checkpoint(&checkpointedClosure, inputClosure)
+
+    // Free the input closure (checkpointed version has its own reference)
+    mlx_closure_free(inputClosure)
+
+    guard status == 0 else {
+        fatalError("mlx_checkpoint failed with status \(status)")
+    }
+
+    // Return a Swift function that invokes the checkpointed closure
+    return { (inputs: [MLXArray]) -> [MLXArray] in
+        // Convert inputs to C vector
+        let inputVec = new_mlx_vector_array(inputs)
+        defer { mlx_vector_array_free(inputVec) }
+
+        // Call the checkpointed closure
+        var outputVec = mlx_vector_array_new()
+        let applyStatus = mlx_closure_apply(&outputVec, checkpointedClosure, inputVec)
+        defer { mlx_vector_array_free(outputVec) }
+
+        guard applyStatus == 0 else {
+            fatalError("mlx_closure_apply failed with status \(applyStatus)")
+        }
+
+        // Convert output back to Swift arrays
+        return mlx_vector_array_values(outputVec)
+    }
+}
+
+/// Convenience overload for single input/output
+public func checkpoint(_ f: @escaping (MLXArray) -> MLXArray) -> (MLXArray) -> MLXArray {
+    let wrapped = checkpoint { (arrays: [MLXArray]) -> [MLXArray] in
+        return [f(arrays[0])]
+    }
+    return { input in wrapped([input])[0] }
+}
+
+// MARK: - Helper functions (copied from MLX internal, needed for C interop)
+
+/// Create a new mlx_vector_array from Swift arrays
+private func new_mlx_vector_array(_ arrays: [MLXArray]) -> mlx_vector_array {
+    withExtendedLifetime(arrays) {
+        mlx_vector_array_new_data(arrays.map { $0.ctx }, arrays.count)
+    }
+}
+
+/// Convert mlx_vector_array to Swift array
+private func mlx_vector_array_values(_ vector_array: mlx_vector_array) -> [MLXArray] {
+    (0 ..< mlx_vector_array_size(vector_array))
+        .map { index in
+            var ctx = mlx_array_new()
+            mlx_vector_array_get(&ctx, vector_array, index)
+            return MLXArray(ctx)
+        }
+}
+
+/// Create a mlx_closure from a Swift function
+private func new_mlx_closure(_ f: @escaping ([MLXArray]) -> [MLXArray]) -> mlx_closure {
+    class ClosureCaptureState {
+        let f: ([MLXArray]) -> [MLXArray]
+        init(_ f: @escaping ([MLXArray]) -> [MLXArray]) {
+            self.f = f
+        }
+    }
+
+    func free(ptr: UnsafeMutableRawPointer?) {
+        Unmanaged<ClosureCaptureState>.fromOpaque(ptr!).release()
+    }
+
+    let payload = Unmanaged.passRetained(ClosureCaptureState(f)).toOpaque()
+
+    func trampoline(
+        resultOut: UnsafeMutablePointer<mlx_vector_array>?,
+        vector_array: mlx_vector_array,
+        payload: UnsafeMutableRawPointer?
+    ) -> Int32 {
+        let state = Unmanaged<ClosureCaptureState>.fromOpaque(payload!).takeUnretainedValue()
+        let arrays = mlx_vector_array_values(vector_array)
+        let result = state.f(arrays)
+
+        if let resultOut {
+            resultOut.pointee = new_mlx_vector_array(result)
+        } else {
+            fatalError("no resultOut pointer")
+        }
+        return 0
+    }
+
+    return mlx_closure_new_func_payload(trampoline, payload, free)
+}