bleam 0.0.8 → 0.0.9

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

@@ -0,0 +1,379 @@
+// Flux2Attention.swift - Joint Attention for Flux.2 Double-Stream Blocks
+// Copyright 2025 Vincent Gourbin
+
+import Foundation
+import MLX
+import MLXNN
+
+/// RMS Normalization (Root Mean Square Layer Normalization)
+/// Optimized using MLXFast.rmsNorm for better performance
+public class FluxRMSNorm: Module, @unchecked Sendable {
+    let dim: Int
+    let eps: Float
+    let weight: MLXArray
+
+    public init(dim: Int, eps: Float = 1e-6) {
+        self.dim = dim
+        self.eps = eps
+        self.weight = MLXArray.ones([dim])
+    }
+
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        // Use optimized MLXFast implementation
+        return MLXFast.rmsNorm(x, weight: weight, eps: eps)
+    }
+}
+
+/// Joint Attention for Double-Stream Blocks
+///
+/// Processes both image hidden states and text encoder hidden states together.
+/// Uses separate Q/K/V projections for each modality, then performs joint attention.
+public class Flux2Attention: Module, @unchecked Sendable {
+    let dim: Int
+    let numHeads: Int
+    let headDim: Int
+
+    // Image hidden states projections (@ModuleInfo for LoRA injection via update)
+    @ModuleInfo var toQ: Linear
+    @ModuleInfo var toK: Linear
+    @ModuleInfo var toV: Linear
+
+    // Text encoder hidden states projections
+    @ModuleInfo var addQProj: Linear
+    @ModuleInfo var addKProj: Linear
+    @ModuleInfo var addVProj: Linear
+
+    // QK normalization (FluxRMSNorm)
+    let normQ: FluxRMSNorm
+    let normK: FluxRMSNorm
+    let normAddedQ: FluxRMSNorm
+    let normAddedK: FluxRMSNorm
+
+    // Output projections
+    @ModuleInfo var toOut: Linear
+    @ModuleInfo var toAddOut: Linear
+
+    /// Initialize Joint Attention
+    /// - Parameters:
+    ///   - dim: Model dimension (6144 for Flux.2)
+    ///   - numHeads: Number of attention heads (48)
+    ///   - headDim: Dimension per head (128)
+    ///   - contextDim: Dimension of text encoder hidden states (same as dim after projection)
+    public init(
+        dim: Int,
+        numHeads: Int,
+        headDim: Int,
+        contextDim: Int? = nil
+    ) {
+        self.dim = dim
+        self.numHeads = numHeads
+        self.headDim = headDim
+
+        let innerDim = numHeads * headDim
+        let ctxDim = contextDim ?? dim
+
+        // Image projections (no bias to match checkpoint)
+        self._toQ.wrappedValue = Linear(dim, innerDim, bias: false)
+        self._toK.wrappedValue = Linear(dim, innerDim, bias: false)
+        self._toV.wrappedValue = Linear(dim, innerDim, bias: false)
+
+        // Text projections (no bias to match checkpoint)
+        self._addQProj.wrappedValue = Linear(ctxDim, innerDim, bias: false)
+        self._addKProj.wrappedValue = Linear(ctxDim, innerDim, bias: false)
+        self._addVProj.wrappedValue = Linear(ctxDim, innerDim, bias: false)
+
+        // QK normalization
+        self.normQ = FluxRMSNorm(dim: headDim)
+        self.normK = FluxRMSNorm(dim: headDim)
+        self.normAddedQ = FluxRMSNorm(dim: headDim)
+        self.normAddedK = FluxRMSNorm(dim: headDim)
+
+        // Output projections (no bias to match checkpoint)
+        self._toOut.wrappedValue = Linear(innerDim, dim, bias: false)
+        self._toAddOut.wrappedValue = Linear(innerDim, ctxDim, bias: false)
+    }
+
+    /// Forward pass for joint attention
+    /// - Parameters:
+    ///   - hiddenStates: Image hidden states [B, S_img, dim]
+    ///   - encoderHiddenStates: Text encoder hidden states [B, S_txt, dim]
+    ///   - rotaryEmb: Optional RoPE embeddings (cos, sin)
+    /// - Returns: Tuple of updated (image hidden states, text hidden states)
+    public func callAsFunction(
+        hiddenStates: MLXArray,
+        encoderHiddenStates: MLXArray,
+        rotaryEmb: (cos: MLXArray, sin: MLXArray)? = nil
+    ) -> (hiddenStates: MLXArray, encoderHiddenStates: MLXArray) {
+        Flux2Debug.verbose("Flux2Attention input - img: \(hiddenStates.shape), txt: \(encoderHiddenStates.shape)")
+
+        let batchSize = hiddenStates.shape[0]
+        let seqLenImg = hiddenStates.shape[1]
+        let seqLenTxt = encoderHiddenStates.shape[1]
+
+        // Project image to Q, K, V
+        var q = toQ(hiddenStates)
+        var k = toK(hiddenStates)
+        let v = toV(hiddenStates)
+        Flux2Debug.verbose("Image Q/K/V projected: q=\(q.shape), k=\(k.shape), v=\(v.shape)")
+
+        // Project text to Q, K, V
+        var addedQ = addQProj(encoderHiddenStates)
+        var addedK = addKProj(encoderHiddenStates)
+        let addedV = addVProj(encoderHiddenStates)
+        Flux2Debug.verbose("Text Q/K/V projected: addedQ=\(addedQ.shape), addedK=\(addedK.shape), addedV=\(addedV.shape)")
+
+        // Reshape for multi-head attention: [B, S, H*D] -> [B, H, S, D]
+        q = reshapeForAttention(q, batchSize: batchSize, seqLen: seqLenImg)
+        k = reshapeForAttention(k, batchSize: batchSize, seqLen: seqLenImg)
+        let vReshaped = reshapeForAttention(v, batchSize: batchSize, seqLen: seqLenImg)
+        Flux2Debug.verbose("Image after reshape: q=\(q.shape), k=\(k.shape), v=\(vReshaped.shape)")
+
+        addedQ = reshapeForAttention(addedQ, batchSize: batchSize, seqLen: seqLenTxt)
+        addedK = reshapeForAttention(addedK, batchSize: batchSize, seqLen: seqLenTxt)
+        let addedVReshaped = reshapeForAttention(addedV, batchSize: batchSize, seqLen: seqLenTxt)
+        Flux2Debug.verbose("Text after reshape: addedQ=\(addedQ.shape), addedK=\(addedK.shape), addedV=\(addedVReshaped.shape)")
+
+        // Apply QK normalization
+        q = normQ(q)
+        k = normK(k)
+        addedQ = normAddedQ(addedQ)
+        addedK = normAddedK(addedK)
+        Flux2Debug.verbose("After QK norm: q=\(q.shape), addedQ=\(addedQ.shape)")
+
+        // Apply RoPE if provided
+        if let rope = rotaryEmb {
+            // Split RoPE for text and image portions
+            // Note: combined IDs are [txtIds, imgIds], so text comes first
+            let txtCos = rope.cos[0..<seqLenTxt]
+            let txtSin = rope.sin[0..<seqLenTxt]
+            let imgCos = rope.cos[seqLenTxt..<(seqLenTxt + seqLenImg)]
+            let imgSin = rope.sin[seqLenTxt..<(seqLenTxt + seqLenImg)]
+
+            Flux2Debug.verbose("RoPE split - txtCos: \(txtCos.shape), imgCos: \(imgCos.shape)")
+
+            (q, k) = applyRoPE(q: q, k: k, cos: imgCos, sin: imgSin)
+            (addedQ, addedK) = applyRoPE(q: addedQ, k: addedK, cos: txtCos, sin: txtSin)
+            Flux2Debug.verbose("After RoPE: q=\(q.shape), addedQ=\(addedQ.shape)")
+        }
+
+        // Concatenate image and text for joint attention
+        let concatQ = concatenated([addedQ, q], axis: 2)  // [B, H, S_txt+S_img, D]
+        let concatK = concatenated([addedK, k], axis: 2)
+        let concatV = concatenated([addedVReshaped, vReshaped], axis: 2)
+        Flux2Debug.verbose("Concatenated for attention: Q=\(concatQ.shape), K=\(concatK.shape), V=\(concatV.shape)")
+
+        // Compute attention using MLXFast scaled dot product attention
+        Flux2Debug.verbose("Calling scaledDotProductAttention...")
+        let attnOutput = MLXFast.scaledDotProductAttention(
+            queries: concatQ,
+            keys: concatK,
+            values: concatV,
+            scale: Float(1.0 / sqrt(Float(headDim))),
+            mask: nil
+        )
+        Flux2Debug.verbose("Attention output: \(attnOutput.shape)")
+
+        // Split back into text and image portions
+        let txtAttnOutput = attnOutput[0..., 0..., 0..<seqLenTxt, 0...]
+        let imgAttnOutput = attnOutput[0..., 0..., seqLenTxt..., 0...]
+        Flux2Debug.verbose("Split attention: txt=\(txtAttnOutput.shape), img=\(imgAttnOutput.shape)")
+
+        // Reshape back: [B, H, S, D] -> [B, S, H*D]
+        let imgOut = reshapeFromAttention(imgAttnOutput, batchSize: batchSize, seqLen: seqLenImg)
+        let txtOut = reshapeFromAttention(txtAttnOutput, batchSize: batchSize, seqLen: seqLenTxt)
+        Flux2Debug.verbose("After reshape from attn: img=\(imgOut.shape), txt=\(txtOut.shape)")
+
+        // Project outputs
+        let hiddenStatesOut = toOut(imgOut)
+        let encoderHiddenStatesOut = toAddOut(txtOut)
+        Flux2Debug.verbose("Final outputs: img=\(hiddenStatesOut.shape), txt=\(encoderHiddenStatesOut.shape)")
+
+        return (hiddenStates: hiddenStatesOut, encoderHiddenStates: encoderHiddenStatesOut)
+    }
+
+    // MARK: - Helper Functions
+
+    /// Forward pass with cached text K/V (for KV cache acceleration)
+    /// - Parameters:
+    ///   - hiddenStates: Image hidden states [B, S_img, dim]
+    ///   - encoderHiddenStates: Text encoder hidden states [B, S_txt, dim] (used only for Q projection)
+    ///   - cachedTextKV: Pre-computed text K and V [B, H, S_txt, D]
+    ///   - rotaryEmb: Optional RoPE embeddings (cos, sin) - only image portion needed
+    /// - Returns: Tuple of updated (image hidden states, text hidden states)
+    public func callAsFunction(
+        hiddenStates: MLXArray,
+        encoderHiddenStates: MLXArray,
+        cachedTextKV: (k: MLXArray, v: MLXArray),
+        rotaryEmb: (cos: MLXArray, sin: MLXArray)? = nil
+    ) -> (hiddenStates: MLXArray, encoderHiddenStates: MLXArray) {
+        let batchSize = hiddenStates.shape[0]
+        let seqLenImg = hiddenStates.shape[1]
+        let seqLenTxt = cachedTextKV.k.shape[2]
+
+        // Project image to Q, K, V
+        var k = toK(hiddenStates)
+        let v = toV(hiddenStates)
+
+        // Project text to Q only (K/V are cached)
+        var addedQ = addQProj(encoderHiddenStates)
+
+        // Reshape for multi-head attention
+        var qImg = toQ(hiddenStates)
+        qImg = reshapeForAttention(qImg, batchSize: batchSize, seqLen: seqLenImg)
+        k = reshapeForAttention(k, batchSize: batchSize, seqLen: seqLenImg)
+        let vReshaped = reshapeForAttention(v, batchSize: batchSize, seqLen: seqLenImg)
+
+        addedQ = reshapeForAttention(addedQ, batchSize: batchSize, seqLen: seqLenTxt)
+
+        // Apply QK normalization
+        qImg = normQ(qImg)
+        k = normK(k)
+        addedQ = normAddedQ(addedQ)
+
+        // Apply RoPE to image Q/K only (text doesn't need it since we're not recomputing its K)
+        if let rope = rotaryEmb {
+            let imgCos = rope.cos[seqLenTxt..<(seqLenTxt + seqLenImg)]
+            let imgSin = rope.sin[seqLenTxt..<(seqLenTxt + seqLenImg)]
+            (qImg, k) = applyRoPE(q: qImg, k: k, cos: imgCos, sin: imgSin)
+        }
+
+        // Concatenate cached text K/V with image K/V for joint attention
+        let concatK = concatenated([cachedTextKV.k, k], axis: 2)
+        let concatV = concatenated([cachedTextKV.v, vReshaped], axis: 2)
+
+        // Compute attention
+        let attnOutput = MLXFast.scaledDotProductAttention(
+            queries: concatenated([addedQ, qImg], axis: 2),
+            keys: concatK,
+            values: concatV,
+            scale: Float(1.0 / sqrt(Float(headDim))),
+            mask: nil
+        )
+
+        // Split back into text and image portions
+        let txtAttnOutput = attnOutput[0..., 0..., 0..<seqLenTxt, 0...]
+        let imgAttnOutput = attnOutput[0..., 0..., seqLenTxt..., 0...]
+
+        // Reshape back
+        let imgOut = reshapeFromAttention(imgAttnOutput, batchSize: batchSize, seqLen: seqLenImg)
+        let txtOut = reshapeFromAttention(txtAttnOutput, batchSize: batchSize, seqLen: seqLenTxt)
+
+        // Project outputs
+        let hiddenStatesOut = toOut(imgOut)
+        let encoderHiddenStatesOut = toAddOut(txtOut)
+
+        return (hiddenStates: hiddenStatesOut, encoderHiddenStates: encoderHiddenStatesOut)
+    }
+
+    /// Compute K and V for text tokens only (for KV cache building)
+    /// - Parameters:
+    ///   - encoderHiddenStates: Normalized text hidden states [B, S_txt, dim]
+    ///   - batchSize: Batch size
+    /// - Returns: (k, v) each [B, H, S_txt, D]
+    public func computeTextKV(
+        encoderHiddenStates: MLXArray,
+        batchSize: Int
+    ) -> (k: MLXArray, v: MLXArray) {
+        let seqLenTxt = encoderHiddenStates.shape[1]
+
+        var addedK = addKProj(encoderHiddenStates)
+        let addedV = addVProj(encoderHiddenStates)
+
+        addedK = reshapeForAttention(addedK, batchSize: batchSize, seqLen: seqLenTxt)
+        let addedVReshaped = reshapeForAttention(addedV, batchSize: batchSize, seqLen: seqLenTxt)
+
+        addedK = normAddedK(addedK)
+
+        return (k: addedK, v: addedVReshaped)
+    }
+
+    /// Self-attention on text tokens only (for KV cache building)
+    /// - Parameter encoderHiddenStates: Normalized text hidden states [B, S_txt, dim]
+    /// - Returns: Text attention output [B, S_txt, dim]
+    public func forwardTextSelfAttention(
+        encoderHiddenStates: MLXArray
+    ) -> MLXArray {
+        let batchSize = encoderHiddenStates.shape[0]
+        let seqLenTxt = encoderHiddenStates.shape[1]
+
+        var addedQ = addQProj(encoderHiddenStates)
+        var addedK = addKProj(encoderHiddenStates)
+        let addedV = addVProj(encoderHiddenStates)
+
+        addedQ = reshapeForAttention(addedQ, batchSize: batchSize, seqLen: seqLenTxt)
+        addedK = reshapeForAttention(addedK, batchSize: batchSize, seqLen: seqLenTxt)
+        let addedVReshaped = reshapeForAttention(addedV, batchSize: batchSize, seqLen: seqLenTxt)
+
+        addedQ = normAddedQ(addedQ)
+        addedK = normAddedK(addedK)
+
+        let attnOutput = MLXFast.scaledDotProductAttention(
+            queries: addedQ,
+            keys: addedK,
+            values: addedVReshaped,
+            scale: Float(1.0 / sqrt(Float(headDim))),
+            mask: nil
+        )
+
+        let txtOut = reshapeFromAttention(attnOutput, batchSize: batchSize, seqLen: seqLenTxt)
+        return toAddOut(txtOut)
+    }
+
+    // MARK: - Helper Functions
+
+    /// Reshape tensor for multi-head attention
+    /// [B, S, H*D] -> [B, H, S, D]
+    private func reshapeForAttention(_ x: MLXArray, batchSize: Int, seqLen: Int) -> MLXArray {
+        x.reshaped([batchSize, seqLen, numHeads, headDim])
+            .transposed(0, 2, 1, 3)
+    }
+
+    /// Reshape tensor from multi-head attention
+    /// [B, H, S, D] -> [B, S, H*D]
+    private func reshapeFromAttention(_ x: MLXArray, batchSize: Int, seqLen: Int) -> MLXArray {
+        x.transposed(0, 2, 1, 3)
+            .reshaped([batchSize, seqLen, numHeads * headDim])
+    }
+
+    /// Apply rotary position embeddings
+    private func applyRoPE(
+        q: MLXArray,
+        k: MLXArray,
+        cos: MLXArray,
+        sin: MLXArray
+    ) -> (MLXArray, MLXArray) {
+        // cos/sin: [S, D] -> [1, 1, S, D]
+        let cosExpanded = cos.expandedDimensions(axes: [0, 1])
+        let sinExpanded = sin.expandedDimensions(axes: [0, 1])
+
+        let qRotated = rotateHalf(q)
+        let kRotated = rotateHalf(k)
+
+        let qOut = q * cosExpanded + qRotated * sinExpanded
+        let kOut = k * cosExpanded + kRotated * sinExpanded
+
+        return (qOut, kOut)
+    }
+
+    /// Rotate features for RoPE (diffusers-style: treat consecutive pairs as real/imag)
+    private func rotateHalf(_ x: MLXArray) -> MLXArray {
+        // x shape: [B, H, S, D]
+        // Diffusers approach: reshape to [B, H, S, D/2, 2], then [-imag, real]
+        let batchSize = x.shape[0]
+        let numHeads = x.shape[1]
+        let seqLen = x.shape[2]
+        let dim = x.shape[3]
+        let halfDim = dim / 2
+
+        // Reshape to [B, H, S, D/2, 2]
+        let xReshaped = x.reshaped([batchSize, numHeads, seqLen, halfDim, 2])
+
+        // Get real and imag parts (consecutive pairs)
+        let xReal = xReshaped[0..., 0..., 0..., 0..., 0]  // [B, H, S, D/2]
+        let xImag = xReshaped[0..., 0..., 0..., 0..., 1]  // [B, H, S, D/2]
+
+        // Create rotated: stack [-imag, real] and flatten
+        let xRotatedStacked = stacked([-xImag, xReal], axis: -1)  // [B, H, S, D/2, 2]
+        return xRotatedStacked.reshaped([batchSize, numHeads, seqLen, dim])  // [B, H, S, D]
+    }
+}

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

@@ -0,0 +1,159 @@
+// Flux2Embeddings.swift - Timestep and guidance embeddings
+// Copyright 2025 Vincent Gourbin
+
+import Foundation
+import MLX
+import MLXNN
+
+/// Sinusoidal timestep embedding projection
+public class Timesteps: Module, @unchecked Sendable {
+    let numChannels: Int
+    let flipSinToCos: Bool
+    let downscaleFreqShift: Float
+    let scale: Float
+
+    public init(
+        numChannels: Int,
+        flipSinToCos: Bool = true,
+        downscaleFreqShift: Float = 0.0,
+        scale: Float = 1.0
+    ) {
+        self.numChannels = numChannels
+        self.flipSinToCos = flipSinToCos
+        self.downscaleFreqShift = downscaleFreqShift
+        self.scale = scale
+    }
+
+    public func callAsFunction(_ timesteps: MLXArray) -> MLXArray {
+        let halfDim = numChannels / 2
+        let exponent = -log(Float(10000)) * MLXArray(0..<halfDim).asType(.float32)
+        let exponentScaled = exponent / (Float(halfDim) - downscaleFreqShift)
+
+        let emb = exp(exponentScaled)
+        let timestepsExpanded = timesteps.expandedDimensions(axis: -1).asType(.float32)
+        let embExpanded = emb.expandedDimensions(axis: 0)
+
+        let sinEmb = sin(timestepsExpanded * embExpanded * scale)
+        let cosEmb = cos(timestepsExpanded * embExpanded * scale)
+
+        if flipSinToCos {
+            return concatenated([cosEmb, sinEmb], axis: -1)
+        } else {
+            return concatenated([sinEmb, cosEmb], axis: -1)
+        }
+    }
+}
+
+/// MLP for embedding timesteps
+public class TimestepEmbedding: Module, @unchecked Sendable {
+    @ModuleInfo var linear1: Linear
+    @ModuleInfo var linear2: Linear
+    let activation: @Sendable (MLXArray) -> MLXArray
+
+    public init(
+        inChannels: Int,
+        timeEmbedDim: Int,
+        activationFn: String = "silu"
+    ) {
+        // No bias to match checkpoint
+        self.linear1 = Linear(inChannels, timeEmbedDim, bias: false)
+        self.linear2 = Linear(timeEmbedDim, timeEmbedDim, bias: false)
+
+        switch activationFn {
+        case "silu", "swish":
+            self.activation = { silu($0) }
+        case "gelu":
+            self.activation = { gelu($0) }
+        case "relu":
+            self.activation = { relu($0) }
+        default:
+            self.activation = { silu($0) }
+        }
+    }
+
+    public func callAsFunction(_ sample: MLXArray) -> MLXArray {
+        var x = linear1(sample)
+        x = activation(x)
+        x = linear2(x)
+        return x
+    }
+}
+
+/// Combined timestep and guidance embeddings for Flux.2
+///
+/// Generates embeddings from timestep and guidance scale values.
+/// Output dimension is the transformer's inner dimension (6144 for Flux.2).
+public class Flux2TimestepGuidanceEmbeddings: Module, @unchecked Sendable {
+    let timeProj: Timesteps
+    let timestepEmbedder: TimestepEmbedding
+    let guidanceEmbedder: TimestepEmbedding?
+    let useGuidanceEmbeds: Bool
+
+    /// Initialize timestep and guidance embeddings
+    /// - Parameters:
+    ///   - embeddingDim: Dimension of sinusoidal projection (256)
+    ///   - timeEmbedDim: Output dimension for timestep MLP (6144)
+    ///   - useGuidanceEmbeds: Whether to include guidance embedding
+    public init(
+        embeddingDim: Int = 256,
+        timeEmbedDim: Int = 6144,
+        useGuidanceEmbeds: Bool = true
+    ) {
+        self.timeProj = Timesteps(numChannels: embeddingDim)
+        self.timestepEmbedder = TimestepEmbedding(
+            inChannels: embeddingDim,
+            timeEmbedDim: timeEmbedDim
+        )
+
+        self.useGuidanceEmbeds = useGuidanceEmbeds
+        if useGuidanceEmbeds {
+            self.guidanceEmbedder = TimestepEmbedding(
+                inChannels: embeddingDim,
+                timeEmbedDim: timeEmbedDim
+            )
+        } else {
+            self.guidanceEmbedder = nil
+        }
+    }
+
+    /// Generate timestep and optional guidance embeddings
+    /// - Parameters:
+    ///   - timestep: Current diffusion timestep [B]
+    ///   - guidance: Guidance scale [B] (optional)
+    /// - Returns: Combined embedding [B, timeEmbedDim]
+    public func callAsFunction(
+        timestep: MLXArray,
+        guidance: MLXArray? = nil
+    ) -> MLXArray {
+        // Project timestep to sinusoidal embedding
+        let tEmb = timeProj(timestep)
+        // Pass through MLP
+        var temb = timestepEmbedder(tEmb)
+
+        // Add guidance embedding if enabled
+        if useGuidanceEmbeds, let guidanceEmbedder = guidanceEmbedder, let g = guidance {
+            let gEmb = timeProj(g)
+            let gemb = guidanceEmbedder(gEmb)
+            temb = temb + gemb
+        }
+
+        return temb
+    }
+}
+
+/// Pooled text embedding projection for conditioning
+public class PooledTextProjection: Module, @unchecked Sendable {
+    @ModuleInfo var linear: Linear
+
+    /// Initialize pooled text projection
+    /// - Parameters:
+    ///   - inputDim: Input dimension from text encoder
+    ///   - outputDim: Output dimension (time embed dim)
+    public init(inputDim: Int = 768, outputDim: Int = 6144) {
+        self._linear = ModuleInfo(wrappedValue: Linear(inputDim, outputDim))
+    }
+
+    public func callAsFunction(_ pooledEmbedding: MLXArray) -> MLXArray {
+        linear(pooledEmbedding)
+    }
+}

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

@@ -0,0 +1,154 @@
+// Flux2FeedForward.swift - SwiGLU FeedForward Network
+// Copyright 2025 Vincent Gourbin
+
+import Foundation
+import MLX
+import MLXNN
+
+/// GEGLU activation: GELU-gated Linear Unit
+/// Used in some transformer variants
+/// Optimized with kernel fusion for the gating operation
+public class GEGLU: Module, @unchecked Sendable {
+    @ModuleInfo var proj: Linear
+    let dim: Int
+    let innerDim: Int
+
+    /// Compiled gating function for kernel fusion
+    nonisolated(unsafe) private static let compiledGate: (MLXArray, MLXArray) -> MLXArray = compile { gate, value in
+        gelu(gate) * value
+    }
+
+    public init(dim: Int, innerDim: Int, bias: Bool = false) {
+        self.dim = dim
+        self.innerDim = innerDim
+        // Projects to 2x inner_dim for gate and value
+        self._proj = ModuleInfo(wrappedValue: Linear(dim, innerDim * 2, bias: bias))
+    }
+
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        let projected = proj(x)
+        // Split into gate and value
+        let chunks = split(projected, parts: 2, axis: -1)
+        let gate = chunks[0]
+        let value = chunks[1]
+        // GEGLU: gelu(gate) * value - compiled for kernel fusion
+        return Self.compiledGate(gate, value)
+    }
+}
+
+/// SwiGLU activation: Swish-gated Linear Unit
+/// Used in Flux.2 transformer feedforward blocks
+/// Optimized with kernel fusion for the gating operation
+public class SwiGLU: Module, @unchecked Sendable {
+    @ModuleInfo var proj: Linear
+    let dim: Int
+    let innerDim: Int
+
+    /// Compiled gating function for kernel fusion
+    nonisolated(unsafe) private static let compiledGate: (MLXArray, MLXArray) -> MLXArray = compile { gate, value in
+        silu(gate) * value
+    }
+
+    public init(dim: Int, innerDim: Int, bias: Bool = false) {
+        self.dim = dim
+        self.innerDim = innerDim
+        // Projects to 2x inner_dim for gate and value
+        self._proj.wrappedValue = Linear(dim, innerDim * 2, bias: bias)
+    }
+
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        let projected = proj(x)
+        // Split into gate and value
+        let chunks = split(projected, parts: 2, axis: -1)
+        let gate = chunks[0]
+        let value = chunks[1]
+        // SwiGLU: silu(gate) * value (silu = swish) - compiled for kernel fusion
+        return Self.compiledGate(gate, value)
+    }
+}
+
+/// FeedForward network for Flux.2 transformer blocks
+///
+/// Architecture: Linear -> SwiGLU -> Linear
+/// - Input: [B, S, dim]
+/// - Output: [B, S, dim]
+public class Flux2FeedForward: Module, @unchecked Sendable {
+    let dim: Int
+    let innerDim: Int
+    @ModuleInfo var activation: SwiGLU
+    @ModuleInfo var linearOut: Linear
+
+    /// Initialize FeedForward block
+    /// - Parameters:
+    ///   - dim: Model dimension (6144 for Flux.2)
+    ///   - innerDim: Inner dimension (typically 4 * dim for standard FFN)
+    ///   - bias: Whether to use bias in linear layers (default false to match checkpoint)
+    public init(
+        dim: Int,
+        innerDim: Int? = nil,
+        bias: Bool = false
+    ) {
+        self.dim = dim
+        // Default inner_dim is 4 * dim, but Flux.2 may use different ratio
+        self.innerDim = innerDim ?? (dim * 4)
+
+        // SwiGLU combines input projection and gating
+        self._activation.wrappedValue = SwiGLU(dim: dim, innerDim: self.innerDim, bias: bias)
+
+        // Output projection
+        self._linearOut.wrappedValue = Linear(self.innerDim, dim, bias: bias)
+    }
+
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        // SwiGLU: projects to 2*innerDim, applies silu gate, outputs innerDim
+        var hidden = activation(x)
+        // Project back to dim
+        hidden = linearOut(hidden)
+        return hidden
+    }
+}
+
+/// FeedForward variant with separate gate and up projections
+/// (Alternative implementation matching some model architectures)
+public class Flux2FeedForwardSplit: Module, @unchecked Sendable {
+    let dim: Int
+    let innerDim: Int
+    @ModuleInfo var linearGate: Linear
+    @ModuleInfo var linearUp: Linear
+    @ModuleInfo var linearDown: Linear
+
+    public init(
+        dim: Int,
+        innerDim: Int? = nil,
+        bias: Bool = false
+    ) {
+        self.dim = dim
+        self.innerDim = innerDim ?? (dim * 4)
+
+        self._linearGate = ModuleInfo(wrappedValue: Linear(dim, self.innerDim, bias: bias))
+        self._linearUp = ModuleInfo(wrappedValue: Linear(dim, self.innerDim, bias: bias))
+        self._linearDown = ModuleInfo(wrappedValue: Linear(self.innerDim, dim, bias: bias))
+    }
+
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        // SwiGLU with separate projections
+        let gate = silu(linearGate(x))
+        let up = linearUp(x)
+        let hidden = gate * up
+        return linearDown(hidden)
+    }
+}
+
+/// Context-specific FeedForward for text encoder hidden states
+/// Used in double-stream blocks for processing text separately
+public class Flux2ContextFeedForward: Module, @unchecked Sendable {
+    let feedForward: Flux2FeedForward
+
+    public init(dim: Int, innerDim: Int? = nil, bias: Bool = false) {
+        self.feedForward = Flux2FeedForward(dim: dim, innerDim: innerDim, bias: bias)
+    }
+
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        feedForward(x)
+    }
+}