bleam 0.0.8 → 0.0.9

package/templates/native/ios/Bleam/AI/Flux2/VAE/ResnetBlock.swift

@@ -0,0 +1,316 @@
+// ResnetBlock.swift - ResNet blocks for VAE
+// Copyright 2025 Vincent Gourbin
+
+import Foundation
+import MLX
+import MLXNN
+
+/// Group Normalization for VAE (NHWC format for MLX)
+public class FluxGroupNorm: Module, @unchecked Sendable {
+    let numGroups: Int
+    let numChannels: Int
+    let eps: Float
+    let weight: MLXArray
+    let bias: MLXArray
+
+    public init(numGroups: Int, numChannels: Int, eps: Float = 1e-6) {
+        self.numGroups = numGroups
+        self.numChannels = numChannels
+        self.eps = eps
+        self.weight = MLXArray.ones([numChannels])
+        self.bias = MLXArray.zeros([numChannels])
+    }
+
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        // x shape: [B, H, W, C] (NHWC format for MLX)
+        let shape = x.shape
+        let B = shape[0]
+        let H = shape[1]
+        let W = shape[2]
+        let C = shape[3]
+
+        // Mixed precision: compute statistics in Float32 for numerical stability
+        let inputDtype = x.dtype
+        let xFloat = x.asType(.float32)
+
+        // Reshape to [B, H, W, G, C/G]
+        let reshaped = xFloat.reshaped([B, H, W, numGroups, C / numGroups])
+
+        // Compute mean and variance per group (over H, W, and C/G within each group)
+        let mean = MLX.mean(reshaped, axes: [1, 2, 4], keepDims: true)
+        let variance = MLX.mean(pow(reshaped - mean, 2), axes: [1, 2, 4], keepDims: true)
+
+        // Normalize in Float32
+        let normalized = (reshaped - mean) / sqrt(variance + eps)
+
+        // Reshape back to [B, H, W, C] and convert to original dtype
+        let out = normalized.reshaped([B, H, W, C]).asType(inputDtype)
+
+        // Apply weight and bias [1, 1, 1, C] for NHWC
+        let weightReshaped = weight.reshaped([1, 1, 1, C])
+        let biasReshaped = bias.reshaped([1, 1, 1, C])
+
+        return out * weightReshaped + biasReshaped
+    }
+}
+
+/// Batch Normalization for VAE latent space (Flux.2 specific) - NHWC format for MLX
+public class BatchNorm2d: Module, @unchecked Sendable {
+    let numFeatures: Int
+    let eps: Float
+    let momentum: Float
+    let trackRunningStats: Bool
+
+    let weight: MLXArray
+    let bias: MLXArray
+    var runningMean: MLXArray
+    var runningVar: MLXArray
+
+    public init(
+        numFeatures: Int,
+        eps: Float = 1e-5,
+        momentum: Float = 0.1,
+        trackRunningStats: Bool = true
+    ) {
+        self.numFeatures = numFeatures
+        self.eps = eps
+        self.momentum = momentum
+        self.trackRunningStats = trackRunningStats
+
+        self.weight = MLXArray.ones([numFeatures])
+        self.bias = MLXArray.zeros([numFeatures])
+        self.runningMean = MLXArray.zeros([numFeatures])
+        self.runningVar = MLXArray.ones([numFeatures])
+    }
+
+    public func callAsFunction(_ x: MLXArray, training: Bool = false) -> MLXArray {
+        // x shape: [B, H, W, C] (NHWC format for MLX)
+        let shape = x.shape
+        let C = shape[3]
+
+        // Mixed precision: compute statistics in Float32 for numerical stability
+        let inputDtype = x.dtype
+        let xFloat = x.asType(.float32)
+
+        if training {
+            // Compute batch statistics (over B, H, W) in Float32
+            let mean = MLX.mean(xFloat, axes: [0, 1, 2])
+            let variance = MLX.mean(pow(xFloat - mean.reshaped([1, 1, 1, C]), 2), axes: [0, 1, 2])
+
+            // Update running stats (keep in Float32)
+            if trackRunningStats {
+                runningMean = (1 - momentum) * runningMean + momentum * mean
+                runningVar = (1 - momentum) * runningVar + momentum * variance
+            }
+
+            // Normalize with batch stats in Float32, then convert back
+            let normalized = (xFloat - mean.reshaped([1, 1, 1, C])) / sqrt(variance.reshaped([1, 1, 1, C]) + eps)
+            let out = normalized.asType(inputDtype)
+            return out * weight.reshaped([1, 1, 1, C]) + bias.reshaped([1, 1, 1, C])
+        } else {
+            // Use running stats - normalize in Float32
+            let runningMeanFloat = runningMean.asType(.float32)
+            let runningVarFloat = runningVar.asType(.float32)
+            let normalized = (xFloat - runningMeanFloat.reshaped([1, 1, 1, C])) / sqrt(runningVarFloat.reshaped([1, 1, 1, C]) + eps)
+            let out = normalized.asType(inputDtype)
+            return out * weight.reshaped([1, 1, 1, C]) + bias.reshaped([1, 1, 1, C])
+        }
+    }
+}
+
+/// ResNet block for VAE encoder/decoder
+public class ResnetBlock2D: Module, @unchecked Sendable {
+    let inChannels: Int
+    let outChannels: Int
+
+    let norm1: FluxGroupNorm
+    let conv1: Conv2d
+    let norm2: FluxGroupNorm
+    let conv2: Conv2d
+    let convShortcut: Conv2d?
+
+    public init(
+        inChannels: Int,
+        outChannels: Int? = nil,
+        numGroups: Int = 32
+    ) {
+        self.inChannels = inChannels
+        self.outChannels = outChannels ?? inChannels
+
+        self.norm1 = FluxGroupNorm(numGroups: numGroups, numChannels: inChannels)
+        self.conv1 = Conv2d(
+            inputChannels: inChannels,
+            outputChannels: self.outChannels,
+            kernelSize: 3,
+            padding: 1
+        )
+
+        self.norm2 = FluxGroupNorm(numGroups: numGroups, numChannels: self.outChannels)
+        self.conv2 = Conv2d(
+            inputChannels: self.outChannels,
+            outputChannels: self.outChannels,
+            kernelSize: 3,
+            padding: 1
+        )
+
+        // Shortcut projection if channels change
+        if inChannels != self.outChannels {
+            self.convShortcut = Conv2d(
+                inputChannels: inChannels,
+                outputChannels: self.outChannels,
+                kernelSize: 1
+            )
+        } else {
+            self.convShortcut = nil
+        }
+    }
+
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        var hidden = norm1(x)
+        hidden = silu(hidden)
+        hidden = conv1(hidden)
+
+        hidden = norm2(hidden)
+        hidden = silu(hidden)
+        hidden = conv2(hidden)
+
+        // Shortcut
+        let shortcut: MLXArray
+        if let convShortcut = convShortcut {
+            shortcut = convShortcut(x)
+        } else {
+            shortcut = x
+        }
+
+        return hidden + shortcut
+    }
+}
+
+/// Downsample block for VAE encoder
+public class Downsample2D: Module, @unchecked Sendable {
+    let conv: Conv2d
+
+    public init(channels: Int, useConv: Bool = true) {
+        if useConv {
+            self.conv = Conv2d(
+                inputChannels: channels,
+                outputChannels: channels,
+                kernelSize: 3,
+                stride: 2
+            )
+        } else {
+            // Average pooling fallback (not typical for VAE)
+            self.conv = Conv2d(
+                inputChannels: channels,
+                outputChannels: channels,
+                kernelSize: 3,
+                stride: 2
+            )
+        }
+    }
+
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        let paddedInput = padded(x, widths: [[0, 0], [0, 1], [0, 1], [0, 0]])
+        return conv(paddedInput)
+    }
+}
+
+/// Upsample block for VAE decoder (NHWC format for MLX)
+public class Upsample2D: Module, @unchecked Sendable {
+    let conv: Conv2d
+
+    public init(channels: Int, useConv: Bool = true) {
+        self.conv = Conv2d(
+            inputChannels: channels,
+            outputChannels: channels,
+            kernelSize: 3,
+            padding: 1
+        )
+    }
+
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        // x shape: [B, H, W, C] (NHWC format for MLX)
+        let shape = x.shape
+        let B = shape[0]
+        let H = shape[1]
+        let W = shape[2]
+        let C = shape[3]
+
+        // Upsample by 2x using nearest neighbor interpolation
+        // The correct approach is to expand dimensions and then reshape:
+        // [B, H, W, C] -> [B, H, 1, W, 1, C] -> broadcast to [B, H, 2, W, 2, C] -> [B, H*2, W*2, C]
+
+        // Reshape to add extra dimensions for broadcasting
+        var upsampled = x.reshaped([B, H, 1, W, 1, C])
+
+        // Tile/broadcast to duplicate each element along H and W
+        // Use concatenation to duplicate along each new axis
+        upsampled = concatenated([upsampled, upsampled], axis: 2)  // [B, H, 2, W, 1, C]
+        upsampled = concatenated([upsampled, upsampled], axis: 4)  // [B, H, 2, W, 2, C]
+
+        // Reshape to merge the duplicated dimensions
+        upsampled = upsampled.reshaped([B, H * 2, W * 2, C])
+
+        return conv(upsampled)
+    }
+}
+
+/// Self-attention block for VAE (mid-block) - NHWC format for MLX
+/// Uses Linear layers (as in HuggingFace checkpoint) instead of Conv2d
+public class AttentionBlock: Module, @unchecked Sendable {
+    let channels: Int
+    let numHeads: Int
+    let headDim: Int
+
+    let groupNorm: FluxGroupNorm
+    let toQ: Linear
+    let toK: Linear
+    let toV: Linear
+    let toOut: Linear
+
+    public init(channels: Int, numHeads: Int = 1, numGroups: Int = 32) {
+        self.channels = channels
+        self.numHeads = numHeads
+        self.headDim = channels / numHeads
+
+        self.groupNorm = FluxGroupNorm(numGroups: numGroups, numChannels: channels)
+        self.toQ = Linear(channels, channels)
+        self.toK = Linear(channels, channels)
+        self.toV = Linear(channels, channels)
+        self.toOut = Linear(channels, channels)
+    }
+
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        // x shape: [B, H, W, C] (NHWC format for MLX)
+        let residual = x
+        let shape = x.shape
+        let B = shape[0]
+        let H = shape[1]
+        let W = shape[2]
+        let C = shape[3]
+
+        // Normalize
+        var hidden = groupNorm(x)
+
+        // Reshape to [B, H*W, C] for Linear projection
+        hidden = hidden.reshaped([B, H * W, C])
+
+        // Project to Q, K, V using Linear layers
+        let q = toQ(hidden)  // [B, HW, C]
+        let k = toK(hidden).transposed(0, 2, 1)  // [B, C, HW]
+        let v = toV(hidden)  // [B, HW, C]
+
+        // Attention
+        let scale = Float(1.0 / sqrt(Float(C)))
+        let attnWeights = softmax(matmul(q, k) * scale, axis: -1)
+        let attnOut = matmul(attnWeights, v)  // [B, HW, C]
+
+        // Project output
+        hidden = toOut(attnOut)
+
+        // Reshape back to [B, H, W, C]
+        hidden = hidden.reshaped([B, H, W, C])
+
+        return hidden + residual
+    }
+}

package/templates/native/ios/Bleam/AI/Flux2/VAE/VAEDecoder.swift

@@ -0,0 +1,120 @@
+// VAEDecoder.swift - VAE Decoder for Flux.2
+// Copyright 2025 Vincent Gourbin
+
+import Foundation
+import MLX
+import MLXNN
+
+/// VAE Decoder for Flux.2
+///
+/// Decodes latent representations back to RGB images.
+/// Architecture:
+/// - Initial conv: latent_channels -> 512
+/// - Mid block with attention
+/// - Up blocks with ResNet blocks and upsampling
+/// - Final conv: 128 -> 3
+public class VAEDecoder: Module, @unchecked Sendable {
+    let config: VAEConfig
+
+    let convIn: Conv2d
+    let midBlock: (resnet1: ResnetBlock2D, attention: AttentionBlock, resnet2: ResnetBlock2D)
+    let upBlocks: [(blocks: [ResnetBlock2D], upsample: Upsample2D?)]
+    let convNormOut: FluxGroupNorm
+    let convOut: Conv2d
+
+    public init(config: VAEConfig = .flux2Dev) {
+        self.config = config
+
+        let blockOutChannels = config.blockOutChannels  // [128, 256, 512, 512]
+        let reversedChannels = blockOutChannels.reversed()  // [512, 512, 256, 128]
+
+        // Initial convolution
+        self.convIn = Conv2d(
+            inputChannels: config.latentChannels,
+            outputChannels: blockOutChannels.last!,
+            kernelSize: 3,
+            padding: 1
+        )
+
+        // Mid block
+        let midChannels = blockOutChannels.last!
+        self.midBlock = (
+            resnet1: ResnetBlock2D(inChannels: midChannels, numGroups: config.normNumGroups),
+            attention: AttentionBlock(channels: midChannels, numGroups: config.normNumGroups),
+            resnet2: ResnetBlock2D(inChannels: midChannels, numGroups: config.normNumGroups)
+        )
+
+        // Up blocks (reversed order)
+        var blocks: [(blocks: [ResnetBlock2D], upsample: Upsample2D?)] = []
+        var prevChannels = midChannels
+
+        for (i, outChannels) in Array(reversedChannels).enumerated() {
+            var resBlocks: [ResnetBlock2D] = []
+
+            // Diffusers adds +1 to layers_per_block for ALL up blocks in the decoder
+            let numLayers = config.layersPerBlock + 1
+
+            for j in 0..<numLayers {
+                let inCh = (j == 0) ? prevChannels : outChannels
+                resBlocks.append(ResnetBlock2D(
+                    inChannels: inCh,
+                    outChannels: outChannels,
+                    numGroups: config.normNumGroups
+                ))
+            }
+            prevChannels = outChannels
+
+            // Upsample except for last block
+            let upsample: Upsample2D?
+            if i < Array(reversedChannels).count - 1 {
+                upsample = Upsample2D(channels: outChannels)
+            } else {
+                upsample = nil
+            }
+
+            blocks.append((blocks: resBlocks, upsample: upsample))
+        }
+        self.upBlocks = blocks
+
+        // Output
+        self.convNormOut = FluxGroupNorm(numGroups: config.normNumGroups, numChannels: blockOutChannels[0])
+        self.convOut = Conv2d(
+            inputChannels: blockOutChannels[0],
+            outputChannels: config.outChannels,
+            kernelSize: 3,
+            padding: 1
+        )
+    }
+
+    public func callAsFunction(_ z: MLXArray) -> MLXArray {
+        // z shape: [B, latent_channels, H/8, W/8] (NCHW from transformer)
+        // Convert to NHWC for MLX Conv2d
+        var hidden = z.transposed(0, 2, 3, 1)  // [B, H/8, W/8, latent_channels]
+
+        // Initial conv
+        hidden = convIn(hidden)
+
+        // Mid block
+        hidden = midBlock.resnet1(hidden)
+        hidden = midBlock.attention(hidden)
+        hidden = midBlock.resnet2(hidden)
+
+        // Up blocks
+        for (resBlocks, upsample) in upBlocks {
+            for resBlock in resBlocks {
+                hidden = resBlock(hidden)
+            }
+            if let us = upsample {
+                hidden = us(hidden)
+            }
+        }
+
+        // Output
+        hidden = convNormOut(hidden)
+        hidden = silu(hidden)
+        hidden = convOut(hidden)
+
+        // Convert back to NCHW for output: [B, H, W, 3] -> [B, 3, H, W]
+        return hidden.transposed(0, 3, 1, 2)
+    }
+}

package/templates/native/ios/Bleam/AI/Flux2/VAE/VAEEncoder.swift

@@ -0,0 +1,136 @@
+// VAEEncoder.swift - VAE Encoder for Flux.2
+// Copyright 2025 Vincent Gourbin
+
+import Foundation
+import MLX
+import MLXNN
+
+/// VAE Encoder for Flux.2
+///
+/// Encodes RGB images to latent space.
+/// Architecture:
+/// - Initial conv: 3 -> 128
+/// - Down blocks with ResNet blocks and downsampling
+/// - Mid block with attention
+/// - Final conv: 512 -> latent_channels * 2 (for mean and variance)
+public class VAEEncoder: Module, @unchecked Sendable {
+    let config: VAEConfig
+
+    let convIn: Conv2d
+    let downBlocks: [(blocks: [ResnetBlock2D], downsample: Downsample2D?)]
+    let midBlock: (resnet1: ResnetBlock2D, attention: AttentionBlock, resnet2: ResnetBlock2D)
+    let convNormOut: FluxGroupNorm
+    let convOut: Conv2d
+
+    public init(config: VAEConfig = .flux2Dev) {
+        self.config = config
+
+        let blockOutChannels = config.blockOutChannels  // [128, 256, 512, 512]
+
+        // Initial convolution
+        self.convIn = Conv2d(
+            inputChannels: config.inChannels,
+            outputChannels: blockOutChannels[0],
+            kernelSize: 3,
+            padding: 1
+        )
+
+        // Down blocks
+        var blocks: [(blocks: [ResnetBlock2D], downsample: Downsample2D?)] = []
+        var prevChannels = blockOutChannels[0]
+
+        for (i, outChannels) in blockOutChannels.enumerated() {
+            var resBlocks: [ResnetBlock2D] = []
+
+            for _ in 0..<config.layersPerBlock {
+                resBlocks.append(ResnetBlock2D(
+                    inChannels: prevChannels,
+                    outChannels: outChannels,
+                    numGroups: config.normNumGroups
+                ))
+                prevChannels = outChannels
+            }
+
+            // Downsample except for last block
+            let downsample: Downsample2D?
+            if i < blockOutChannels.count - 1 {
+                downsample = Downsample2D(channels: outChannels)
+            } else {
+                downsample = nil
+            }
+
+            blocks.append((blocks: resBlocks, downsample: downsample))
+        }
+        self.downBlocks = blocks
+
+        // Mid block
+        let midChannels = blockOutChannels.last!
+        self.midBlock = (
+            resnet1: ResnetBlock2D(inChannels: midChannels, numGroups: config.normNumGroups),
+            attention: AttentionBlock(channels: midChannels, numGroups: config.normNumGroups),
+            resnet2: ResnetBlock2D(inChannels: midChannels, numGroups: config.normNumGroups)
+        )
+
+        // Output
+        self.convNormOut = FluxGroupNorm(numGroups: config.normNumGroups, numChannels: midChannels)
+        // Output channels: 2 * latent_channels for mean and log_variance
+        self.convOut = Conv2d(
+            inputChannels: midChannels,
+            outputChannels: config.latentChannels * 2,
+            kernelSize: 3,
+            padding: 1
+        )
+    }
+
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        // x shape: [B, 3, H, W] (NCHW from input)
+        // Convert to NHWC for MLX Conv2d
+        var hidden = x.transposed(0, 2, 3, 1)  // [B, H, W, 3]
+
+        // Initial conv
+        hidden = convIn(hidden)
+
+        // Down blocks
+        for (resBlocks, downsample) in downBlocks {
+            for resBlock in resBlocks {
+                hidden = resBlock(hidden)
+            }
+            if let ds = downsample {
+                hidden = ds(hidden)
+            }
+        }
+
+        // Mid block
+        hidden = midBlock.resnet1(hidden)
+        hidden = midBlock.attention(hidden)
+        hidden = midBlock.resnet2(hidden)
+
+        // Output
+        hidden = convNormOut(hidden)
+        hidden = silu(hidden)
+        hidden = convOut(hidden)
+
+        // Convert back to NCHW: [B, H/8, W/8, latent_channels*2] -> [B, latent_channels*2, H/8, W/8]
+        return hidden.transposed(0, 3, 1, 2)
+    }
+
+    /// Encode and sample from the latent distribution
+    /// - Parameter x: Input image [B, 3, H, W] (NCHW format)
+    /// - Returns: Sampled latent [B, latent_channels, H/8, W/8] (NCHW format)
+    public func encode(_ x: MLXArray, samplePosterior: Bool = true) -> MLXArray {
+        let h = self.callAsFunction(x)  // Output is NCHW
+
+        // Split into mean and log_variance (NCHW format)
+        let mean = h[0..., 0..<config.latentChannels, 0..., 0...]
+        let logVar = h[0..., config.latentChannels..., 0..., 0...]
+
+        if samplePosterior {
+            // Reparameterization trick
+            let std = exp(0.5 * logVar)
+            let noise = MLXRandom.normal(mean.shape)
+            return mean + std * noise
+        } else {
+            return mean
+        }
+    }
+}