bleam 0.0.7 → 0.0.9

package/templates/native/ios/Bleam/AI/Flux2/Pipeline/LatentUtils.swift

@@ -0,0 +1,191 @@
+import Foundation
+import MLX
+
+public enum LatentUtils {
+    public static func generatePatchifiedLatents(
+        batchSize: Int = 1,
+        height: Int,
+        width: Int,
+        latentChannels: Int = 32,
+        patchSize: Int = 2,
+        seed: UInt64? = nil
+    ) -> MLXArray {
+        if let seed {
+            MLXRandom.seed(seed)
+        }
+
+        let patchifiedChannels = latentChannels * patchSize * patchSize
+        let patchifiedHeight = height / (8 * patchSize)
+        let patchifiedWidth = width / (8 * patchSize)
+
+        return MLXRandom.normal([batchSize, patchifiedChannels, patchifiedHeight, patchifiedWidth])
+    }
+
+    public static func packPatchifiedToSequence(_ patchified: MLXArray) -> MLXArray {
+        let shape = patchified.shape
+        let batch = shape[0]
+        let channels = shape[1]
+        let height = shape[2]
+        let width = shape[3]
+
+        return patchified
+            .transposed(0, 2, 3, 1)
+            .reshaped([batch, height * width, channels])
+    }
+
+    public static func unpackSequenceToPatchified(
+        _ sequence: MLXArray,
+        height: Int,
+        width: Int
+    ) -> MLXArray {
+        let shape = sequence.shape
+        let batch = shape[0]
+        let channels = shape[2]
+        let patchifiedHeight = height / 16
+        let patchifiedWidth = width / 16
+
+        return sequence
+            .reshaped([batch, patchifiedHeight, patchifiedWidth, channels])
+            .transposed(0, 3, 1, 2)
+    }
+
+    public static func unpatchifyLatents(
+        _ patchified: MLXArray,
+        latentChannels: Int = 32,
+        patchSize: Int = 2
+    ) -> MLXArray {
+        let shape = patchified.shape
+        let batch = shape[0]
+        let patchifiedHeight = shape[2]
+        let patchifiedWidth = shape[3]
+
+        return patchified
+            .reshaped([batch, latentChannels, patchSize, patchSize, patchifiedHeight, patchifiedWidth])
+            .transposed(0, 1, 4, 2, 5, 3)
+            .reshaped([batch, latentChannels, patchifiedHeight * patchSize, patchifiedWidth * patchSize])
+    }
+
+    public static func packLatentsToPatchified(
+        _ latents: MLXArray,
+        patchSize: Int = 2
+    ) -> MLXArray {
+        let shape = latents.shape
+        let batch = shape[0]
+        let channels = shape[1]
+        let height = shape[2]
+        let width = shape[3]
+        let patchifiedHeight = height / patchSize
+        let patchifiedWidth = width / patchSize
+        let patchifiedChannels = channels * patchSize * patchSize
+
+        return latents
+            .reshaped([batch, channels, patchifiedHeight, patchSize, patchifiedWidth, patchSize])
+            .transposed(0, 2, 4, 1, 3, 5)
+            .reshaped([batch, patchifiedHeight, patchifiedWidth, patchifiedChannels])
+            .transposed(0, 3, 1, 2)
+    }
+
+    public static func generateImagePositionIDs(
+        height: Int,
+        width: Int,
+        patchSize: Int = 2
+    ) -> MLXArray {
+        let latentHeight = height / 8
+        let latentWidth = width / 8
+        let patchesHigh = latentHeight / patchSize
+        let patchesWide = latentWidth / patchSize
+        let patchCount = patchesHigh * patchesWide
+
+        let hIndices = MLXArray.arange(patchesHigh, dtype: .int32)
+        let wIndices = MLXArray.arange(patchesWide, dtype: .int32)
+        let hGrid = MLX.broadcast(hIndices.expandedDimensions(axis: 1), to: [patchesHigh, patchesWide])
+        let wGrid = MLX.broadcast(wIndices.expandedDimensions(axis: 0), to: [patchesHigh, patchesWide])
+        let zeros = MLXArray.zeros([patchCount], dtype: .int32)
+
+        return MLX.stacked([
+            zeros,
+            hGrid.reshaped([patchCount]),
+            wGrid.reshaped([patchCount]),
+            zeros
+        ], axis: 1)
+    }
+
+    public static func generateTextPositionIDs(length: Int) -> MLXArray {
+        let tokenPositions = MLXArray.arange(length, dtype: .int32)
+        let zeros = MLXArray.zeros([length], dtype: .int32)
+
+        return MLX.stacked([zeros, zeros, zeros, tokenPositions], axis: 1)
+    }
+
+    public static func generateReferenceImagePositionIDs(
+        latentHeights: [Int],
+        latentWidths: [Int],
+        scale: Int = 10
+    ) -> MLXArray {
+        var positions: [Int32] = []
+
+        for (imageIndex, pair) in zip(latentHeights, latentWidths).enumerated() {
+            let time = Int32(scale + scale * imageIndex)
+            let (height, width) = pair
+            for y in 0..<height {
+                for x in 0..<width {
+                    positions.append(contentsOf: [time, Int32(y), Int32(x), 0])
+                }
+            }
+        }
+
+        let totalLength = zip(latentHeights, latentWidths).reduce(0) { total, pair in
+            total + pair.0 * pair.1
+        }
+        return MLXArray(positions).reshaped([totalLength, 4])
+    }
+
+    public static func combinePositionIDs(
+        textLength: Int,
+        height: Int,
+        width: Int,
+        patchSize: Int = 2
+    ) -> (textIds: MLXArray, imageIds: MLXArray, combinedIds: MLXArray) {
+        let textIds = generateTextPositionIDs(length: textLength)
+        let imageIds = generateImagePositionIDs(height: height, width: width, patchSize: patchSize)
+        return (textIds, imageIds, concatenated([textIds, imageIds], axis: 0))
+    }
+
+    public static func validateDimensions(
+        height: Int,
+        width: Int,
+        patchSize: Int = 2
+    ) -> (height: Int, width: Int) {
+        let factor = 8 * patchSize
+        return (
+            height: ((height + factor - 1) / factor) * factor,
+            width: ((width + factor - 1) / factor) * factor
+        )
+    }
+
+    public static func denormalizeLatentsWithBatchNorm(
+        _ latents: MLXArray,
+        runningMean: MLXArray,
+        runningVar: MLXArray,
+        eps: Float = 1e-4
+    ) -> MLXArray {
+        let channels = runningMean.shape[0]
+        let mean = runningMean.reshaped([1, channels, 1, 1])
+        let std = sqrt(runningVar.reshaped([1, channels, 1, 1]) + eps)
+
+        return latents * std + mean
+    }
+
+    public static func normalizeLatentsWithBatchNorm(
+        _ latents: MLXArray,
+        runningMean: MLXArray,
+        runningVar: MLXArray,
+        eps: Float = 1e-4
+    ) -> MLXArray {
+        let channels = runningMean.shape[0]
+        let mean = runningMean.reshaped([1, channels, 1, 1])
+        let std = sqrt(runningVar.reshaped([1, channels, 1, 1]) + eps)
+
+        return (latents - mean) / std
+    }
+}

package/templates/native/ios/Bleam/AI/Flux2/Scheduler/FlowMatchEulerScheduler.swift

@@ -0,0 +1,279 @@
+// FlowMatchEulerScheduler.swift - Flow Matching Euler Scheduler
+// Copyright 2025 Vincent Gourbin
+
+import Foundation
+import MLX
+
+/// Compute empirical mu for Flux.2 time shifting
+/// Ported from diffusers: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/flux2/pipeline_flux2.py
+public func computeEmpiricalMu(imageSeqLen: Int, numSteps: Int) -> Float {
+    let a1: Float = 8.73809524e-05
+    let b1: Float = 1.89833333
+    let a2: Float = 0.00016927
+    let b2: Float = 0.45666666
+
+    if imageSeqLen > 4300 {
+        let mu = a2 * Float(imageSeqLen) + b2
+        return mu
+    }
+
+    let m_200 = a2 * Float(imageSeqLen) + b2
+    let m_10 = a1 * Float(imageSeqLen) + b1
+
+    let a = (m_200 - m_10) / 190.0
+    let b = m_200 - 200.0 * a
+    let mu = a * Float(numSteps) + b
+
+    return mu
+}
+
+/// Flow Matching Euler Discrete Scheduler for Flux.2
+///
+/// Implements the flow matching ODE solver with Euler steps.
+/// Based on diffusers FlowMatchEulerDiscreteScheduler
+public class FlowMatchEulerScheduler: @unchecked Sendable {
+    /// Number of training timesteps
+    public let numTrainTimesteps: Int
+
+    /// Shift parameter for timestep scheduling
+    public let shift: Float
+
+    /// Current inference timesteps (in [0, 1000] range for compatibility)
+    public private(set) var timesteps: [Float] = []
+
+    /// Current sigmas (noise levels in [0, 1] range)
+    public private(set) var sigmas: [Float] = []
+
+    /// Step index during sampling
+    public private(set) var stepIndex: Int = 0
+
+    public init(
+        numTrainTimesteps: Int = 1000,
+        shift: Float = 1.0
+    ) {
+        self.numTrainTimesteps = numTrainTimesteps
+        self.shift = shift
+    }
+
+    /// Set timesteps for inference with Flux.2 specific scheduling
+    /// - Parameters:
+    ///   - numInferenceSteps: Number of denoising steps
+    ///   - imageSeqLen: Length of image sequence (for mu calculation)
+    ///   - strength: Denoising strength for I2I (1.0 = full denoise, 0.5 = start at 50%)
+    /// - Returns: The initial timestep index (0 for T2I, >0 for I2I with strength < 1.0)
+    @discardableResult
+    public func setTimesteps(numInferenceSteps: Int, imageSeqLen: Int? = nil, strength: Float = 1.0) -> Int {
+        // Compute mu based on image sequence length (Flux.2 specific), unless a
+        // model profile provided an explicit shift override.
+        let mu: Float
+        if shift != 1.0 {
+            mu = shift
+        } else if let seqLen = imageSeqLen {
+            mu = computeEmpiricalMu(imageSeqLen: seqLen, numSteps: numInferenceSteps)
+        } else {
+            // Default mu for 1024x1024 (4096 latent patches)
+            mu = computeEmpiricalMu(imageSeqLen: 4096, numSteps: numInferenceSteps)
+        }
+
+        // Generate sigmas: linspace(1.0, 1/num_steps, num_steps)
+        // This is the Flux.2 specific sigma schedule
+        var allSigmas: [Float] = []
+        for i in 0..<numInferenceSteps {
+            let sigma = 1.0 - Float(i) / Float(numInferenceSteps)
+            allSigmas.append(sigma)
+        }
+
+        // Apply time shifting with mu (exponential shift)
+        allSigmas = allSigmas.map { sigma in
+            timeShift(mu: mu, sigma: 1.0, t: sigma)
+        }
+
+        // Append terminal sigma (0)
+        allSigmas.append(0.0)
+
+        // For I2I with strength < 1.0, skip early steps
+        // strength = 1.0 → start from step 0 (full denoising)
+        // strength = 0.5 → start from step numSteps/2 (50% denoising)
+        // strength = 0.1 → start from step numSteps*0.9 (10% denoising)
+        let clampedStrength = max(0.01, min(1.0, strength))
+        let initTimestepIndex = numInferenceSteps - Int(Float(numInferenceSteps) * clampedStrength)
+        let tStart = max(0, initTimestepIndex)
+
+        // Slice sigmas from the starting point
+        self.sigmas = Array(allSigmas[tStart...])
+
+        // Timesteps = sigmas * num_train_timesteps (for compatibility)
+        self.timesteps = self.sigmas.map { $0 * Float(numTrainTimesteps) }
+
+        self.stepIndex = 0
+
+        let effectiveSteps = sigmas.count - 1
+        Flux2Debug.log("Scheduler set: \(effectiveSteps) effective steps (strength=\(clampedStrength), mu=\(mu), shift=\(shift))")
+        if strength < 1.0 {
+            Flux2Debug.log("I2I mode: starting from timestep \(tStart) (skipping \(tStart) steps)")
+        }
+        Flux2Debug.verbose("Sigmas: \(sigmas.prefix(5))... to \(sigmas.suffix(2))")
+
+        return tStart
+    }
+
+    /// Get the initial sigma for noise injection (used by I2I)
+    public var initialSigma: Float {
+        sigmas.first ?? 1.0
+    }
+
+    /// Time shift function (exponential) - matches diffusers _time_shift_exponential
+    private func timeShift(mu: Float, sigma: Float, t: Float) -> Float {
+        // exponential time shift: exp(mu) / (exp(mu) + (1/t - 1)^sigma)
+        let expMu = exp(mu)
+        let denominator = expMu + pow(1.0 / t - 1.0, sigma)
+        return expMu / denominator
+    }
+
+    /// Perform one Euler step
+    /// - Parameters:
+    ///   - modelOutput: Predicted velocity from transformer
+    ///   - timestep: Current timestep (not used, we track via stepIndex)
+    ///   - sample: Current noisy sample
+    /// - Returns: Updated sample for next step
+    public func step(
+        modelOutput: MLXArray,
+        timestep: Float,
+        sample: MLXArray
+    ) -> MLXArray {
+        guard stepIndex < sigmas.count - 1 else {
+            return sample
+        }
+
+        let sigma = sigmas[stepIndex]
+        let sigmaNext = sigmas[stepIndex + 1]
+
+        // Euler step: x_{t-1} = x_t + (sigma_next - sigma) * v_t
+        // where v_t is the velocity (model output)
+        let dt = sigmaNext - sigma
+        let nextSample = sample + dt * modelOutput
+
+        stepIndex += 1
+
+        return nextSample
+    }
+
+    /// Get the current sigma for this step
+    public var currentSigma: Float {
+        guard stepIndex < sigmas.count else { return 0 }
+        return sigmas[stepIndex]
+    }
+
+    /// Get initial noise scale
+    public func initNoiseSigma(for timestep: Float) -> Float {
+        // For flow matching starting at t=1, the initial noise scale is 1
+        return 1.0
+    }
+
+    /// Scale model input (identity for flow matching)
+    public func scaleModelInput(_ sample: MLXArray, timestep: Float) -> MLXArray {
+        sample
+    }
+
+    /// Add noise to latents for a given timestep
+    public func addNoise(
+        originalSamples: MLXArray,
+        noise: MLXArray,
+        timestep: Float
+    ) -> MLXArray {
+        // For flow matching: x_t = (1 - t) * x_0 + t * noise
+        // where t is sigma (in [0, 1] range)
+        let sigma = timestep / Float(numTrainTimesteps)
+        let t = MLXArray(sigma)
+        return (1 - t) * originalSamples + t * noise
+    }
+
+    /// Scale noise for image-to-image generation (flow matching)
+    /// This is the core operation for I2I: mix the encoded image with noise based on sigma
+    /// - Parameters:
+    ///   - sample: Original image latents (encoded by VAE)
+    ///   - sigma: Current sigma/noise level (from the time-shifted schedule, in [0, 1] range)
+    ///   - noise: Random noise tensor
+    /// - Returns: Noisy latents: (1 - sigma) * sample + sigma * noise
+    public func scaleNoise(
+        sample: MLXArray,
+        sigma: Float,
+        noise: MLXArray
+    ) -> MLXArray {
+        // For flow matching: x_t = (1 - t) * x_0 + t * noise
+        // sigma is already in [0, 1] range after time shifting
+        let t = MLXArray(sigma)
+        return (1 - t) * sample + t * noise
+    }
+
+    /// Get velocity target for training
+    public func getVelocity(
+        sample: MLXArray,
+        noise: MLXArray,
+        timestep: Float
+    ) -> MLXArray {
+        // Velocity target: v = noise - sample
+        noise - sample
+    }
+
+    /// Reset scheduler state
+    public func reset() {
+        stepIndex = 0
+    }
+
+    /// Set timesteps with custom sigmas (for Turbo LoRAs)
+    ///
+    /// This bypasses the normal mu/timeshift computation and directly uses
+    /// the provided sigma schedule. Used by specialized LoRAs like Flux.2 Turbo
+    /// that require pre-computed noise schedules.
+    ///
+    /// The custom sigmas define the noise levels to denoise through. If the final sigma
+    /// is not 0.0, a terminal sigma of 0.0 is automatically appended to ensure proper
+    /// denoising completion. This matches the standard scheduler behavior where N sigmas
+    /// yield N-1 intervals, so N custom sigmas + terminal 0.0 = N denoising steps.
+    ///
+    /// - Parameter customSigmas: Pre-computed sigma schedule.
+    ///   Should start near 1.0 and end with the second-to-last noise level.
+    ///   Example for 8-step Turbo: [1.0, 0.6509, 0.4374, 0.2932, 0.1893, 0.1108, 0.0495, 0.00031]
+    ///   This will produce 8 denoising steps (with terminal 0.0 appended).
+    public func setCustomSigmas(_ customSigmas: [Float]) {
+        guard !customSigmas.isEmpty else {
+            Flux2Debug.log("[Scheduler] Warning: Empty custom sigmas provided, ignoring")
+            return
+        }
+
+        // Append terminal sigma 0.0 if not already present
+        // This ensures N custom sigmas produce N denoising steps
+        // Even if the last sigma is very small (like 0.00031), we need to add 0.0
+        // to complete the final denoising step
+        var sigmasWithTerminal = customSigmas
+        let lastSigma = customSigmas.last ?? 0.0
+        if lastSigma != 0.0 {  // If last sigma is not exactly 0.0, append terminal
+            sigmasWithTerminal.append(0.0)
+            Flux2Debug.log("[Scheduler] Appended terminal sigma 0.0 (last was \(lastSigma))")
+        }
+
+        self.sigmas = sigmasWithTerminal
+        self.timesteps = sigmasWithTerminal.map { $0 * Float(numTrainTimesteps) }
+        self.stepIndex = 0
+
+        let effectiveSteps = sigmas.count - 1
+        Flux2Debug.log("[Scheduler] Set custom sigmas: \(effectiveSteps) effective steps from \(customSigmas.count) input sigmas")
+        Flux2Debug.verbose("[Scheduler] Sigmas: \(sigmas)")
+    }
+}
+
+// MARK: - Progress Tracking
+
+extension FlowMatchEulerScheduler {
+    /// Current progress (0.0 to 1.0)
+    public var progress: Float {
+        guard !timesteps.isEmpty else { return 0 }
+        return Float(stepIndex) / Float(timesteps.count - 1)
+    }
+
+    /// Remaining steps
+    public var remainingSteps: Int {
+        max(0, timesteps.count - 1 - stepIndex)
+    }
+}