bleam 0.0.8 → 0.0.9

package/templates/native/ios/Bleam/AI/Flux2/Text/KleinTextEncoder.swift

@@ -0,0 +1,380 @@
+import Foundation
+import MLX
+import MLXNN
+
+public struct Qwen3TextConfig: Codable, Sendable {
+    public let vocabSize: Int
+    public let hiddenSize: Int
+    public let intermediateSize: Int
+    public let numHiddenLayers: Int
+    public let numAttentionHeads: Int
+    public let numKeyValueHeads: Int
+    public let maxPositionEmbeddings: Int
+    public let rmsNormEps: Float
+    public let ropeTheta: Float
+    public let tieWordEmbeddings: Bool
+    public let hiddenAct: String
+    public let attentionBias: Bool
+    public let attentionDropout: Float
+    public let headDim: Int
+
+    enum CodingKeys: String, CodingKey {
+        case vocabSize = "vocab_size"
+        case hiddenSize = "hidden_size"
+        case intermediateSize = "intermediate_size"
+        case numHiddenLayers = "num_hidden_layers"
+        case numAttentionHeads = "num_attention_heads"
+        case numKeyValueHeads = "num_key_value_heads"
+        case maxPositionEmbeddings = "max_position_embeddings"
+        case rmsNormEps = "rms_norm_eps"
+        case ropeTheta = "rope_theta"
+        case tieWordEmbeddings = "tie_word_embeddings"
+        case hiddenAct = "hidden_act"
+        case attentionBias = "attention_bias"
+        case attentionDropout = "attention_dropout"
+        case headDim = "head_dim"
+    }
+
+    public init(from decoder: Decoder) throws {
+        let container = try decoder.container(keyedBy: CodingKeys.self)
+        vocabSize = try container.decode(Int.self, forKey: .vocabSize)
+        hiddenSize = try container.decode(Int.self, forKey: .hiddenSize)
+        intermediateSize = try container.decode(Int.self, forKey: .intermediateSize)
+        numHiddenLayers = try container.decode(Int.self, forKey: .numHiddenLayers)
+        numAttentionHeads = try container.decode(Int.self, forKey: .numAttentionHeads)
+        numKeyValueHeads = try container.decode(Int.self, forKey: .numKeyValueHeads)
+        maxPositionEmbeddings = try container.decode(Int.self, forKey: .maxPositionEmbeddings)
+        rmsNormEps = try container.decodeIfPresent(Float.self, forKey: .rmsNormEps) ?? 1e-6
+        ropeTheta = try container.decodeIfPresent(Float.self, forKey: .ropeTheta) ?? 1_000_000.0
+        tieWordEmbeddings = try container.decodeIfPresent(Bool.self, forKey: .tieWordEmbeddings) ?? true
+        hiddenAct = try container.decodeIfPresent(String.self, forKey: .hiddenAct) ?? "silu"
+        attentionBias = try container.decodeIfPresent(Bool.self, forKey: .attentionBias) ?? false
+        attentionDropout = try container.decodeIfPresent(Float.self, forKey: .attentionDropout) ?? 0.0
+        headDim = try container.decodeIfPresent(Int.self, forKey: .headDim) ?? (hiddenSize / numAttentionHeads)
+    }
+
+    public static func load(from url: URL) throws -> Qwen3TextConfig {
+        try JSONDecoder().decode(Qwen3TextConfig.self, from: Data(contentsOf: url))
+    }
+}
+
+private struct Qwen3QuantizationConfig: Codable {
+    let groupSize: Int
+    let bits: Int
+
+    enum CodingKeys: String, CodingKey {
+        case groupSize = "group_size"
+        case bits
+    }
+}
+
+private struct Qwen3ModelConfigWithQuantization: Codable {
+    let quantization: Qwen3QuantizationConfig?
+}
+
+private struct TextEncoderMemoryConfig: Sendable {
+    let evalFrequency: Int
+    let clearCacheOnEval: Bool
+
+    static let aggressive = TextEncoderMemoryConfig(evalFrequency: 4, clearCacheOnEval: true)
+}
+
+public class TextRMSNorm: Module, UnaryLayer {
+    var weight: MLXArray
+    let eps: Float
+
+    public init(dimensions: Int, eps: Float = 1e-5) {
+        self.weight = MLXArray.ones([dimensions])
+        self.eps = eps
+        super.init()
+    }
+
+    public func callAsFunction(_ x: MLXArray) -> MLXArray {
+        let variance = mean(x * x, axis: -1, keepDims: true)
+        return weight * x * rsqrt(variance + MLXArray([eps]))
+    }
+}
+
+private final class Qwen3RoPE: Module {
+    let dimensions: Int
+    let base: Float
+
+    init(dimensions: Int, base: Float) {
+        self.dimensions = dimensions
+        self.base = base
+        super.init()
+    }
+
+    func callAsFunction(_ x: MLXArray) -> MLXArray {
+        MLXFast.RoPE(x, dimensions: dimensions, traditional: false, base: base, scale: 1.0, offset: 0)
+    }
+}
+
+private final class Qwen3Attention: Module {
+    let hiddenSize: Int
+    let numHeads: Int
+    let numKVHeads: Int
+    let headDim: Int
+    let scale: Float
+
+    @ModuleInfo var q_proj: Linear
+    @ModuleInfo var k_proj: Linear
+    @ModuleInfo var v_proj: Linear
+    @ModuleInfo var o_proj: Linear
+    @ModuleInfo var q_norm: TextRMSNorm
+    @ModuleInfo var k_norm: TextRMSNorm
+    let rope: Qwen3RoPE
+
+    init(config: Qwen3TextConfig) {
+        hiddenSize = config.hiddenSize
+        numHeads = config.numAttentionHeads
+        numKVHeads = config.numKeyValueHeads
+        headDim = config.headDim
+        scale = 1.0 / sqrt(Float(headDim))
+        _q_proj = ModuleInfo(wrappedValue: Linear(hiddenSize, numHeads * headDim, bias: config.attentionBias))
+        _k_proj = ModuleInfo(wrappedValue: Linear(hiddenSize, numKVHeads * headDim, bias: config.attentionBias))
+        _v_proj = ModuleInfo(wrappedValue: Linear(hiddenSize, numKVHeads * headDim, bias: config.attentionBias))
+        _o_proj = ModuleInfo(wrappedValue: Linear(numHeads * headDim, hiddenSize, bias: config.attentionBias))
+        _q_norm = ModuleInfo(wrappedValue: TextRMSNorm(dimensions: headDim, eps: config.rmsNormEps))
+        _k_norm = ModuleInfo(wrappedValue: TextRMSNorm(dimensions: headDim, eps: config.rmsNormEps))
+        rope = Qwen3RoPE(dimensions: headDim, base: config.ropeTheta)
+        super.init()
+    }
+
+    func callAsFunction(_ hiddenStates: MLXArray, mask: MLXArray?) -> MLXArray {
+        let batchSize = hiddenStates.shape[0]
+        let seqLen = hiddenStates.shape[1]
+
+        var queries = q_proj(hiddenStates).reshaped([batchSize, seqLen, numHeads, headDim])
+        var keys = k_proj(hiddenStates).reshaped([batchSize, seqLen, numKVHeads, headDim])
+        var values = v_proj(hiddenStates).reshaped([batchSize, seqLen, numKVHeads, headDim])
+
+        queries = q_norm(queries).transposed(0, 2, 1, 3)
+        keys = k_norm(keys).transposed(0, 2, 1, 3)
+        values = values.transposed(0, 2, 1, 3)
+
+        queries = rope(queries)
+        keys = rope(keys)
+
+        let repeatFactor = numHeads / numKVHeads
+        if repeatFactor > 1 {
+            keys = MLX.broadcast(keys.expandedDimensions(axis: 2), to: [batchSize, numKVHeads, repeatFactor, seqLen, headDim])
+                .reshaped([batchSize, numHeads, seqLen, headDim])
+            values = MLX.broadcast(values.expandedDimensions(axis: 2), to: [batchSize, numKVHeads, repeatFactor, seqLen, headDim])
+                .reshaped([batchSize, numHeads, seqLen, headDim])
+        }
+
+        let output = MLXFast.scaledDotProductAttention(
+            queries: queries,
+            keys: keys,
+            values: values,
+            scale: scale,
+            mask: mask
+        )
+
+        return o_proj(output.transposed(0, 2, 1, 3).reshaped([batchSize, seqLen, numHeads * headDim]))
+    }
+}
+
+private final class Qwen3MLP: Module, @unchecked Sendable {
+    @ModuleInfo var gate_proj: Linear
+    @ModuleInfo var up_proj: Linear
+    @ModuleInfo var down_proj: Linear
+
+    nonisolated(unsafe) private static let compiledGate: (MLXArray, MLXArray) -> MLXArray = compile { gate, up in
+        silu(gate) * up
+    }
+
+    init(config: Qwen3TextConfig) {
+        _gate_proj = ModuleInfo(wrappedValue: Linear(config.hiddenSize, config.intermediateSize, bias: false))
+        _up_proj = ModuleInfo(wrappedValue: Linear(config.hiddenSize, config.intermediateSize, bias: false))
+        _down_proj = ModuleInfo(wrappedValue: Linear(config.intermediateSize, config.hiddenSize, bias: false))
+        super.init()
+    }
+
+    func callAsFunction(_ x: MLXArray) -> MLXArray {
+        down_proj(Self.compiledGate(gate_proj(x), up_proj(x)))
+    }
+}
+
+private final class Qwen3DecoderLayer: Module {
+    let self_attn: Qwen3Attention
+    let mlp: Qwen3MLP
+    let input_layernorm: TextRMSNorm
+    let post_attention_layernorm: TextRMSNorm
+
+    init(config: Qwen3TextConfig) {
+        self_attn = Qwen3Attention(config: config)
+        mlp = Qwen3MLP(config: config)
+        input_layernorm = TextRMSNorm(dimensions: config.hiddenSize, eps: config.rmsNormEps)
+        post_attention_layernorm = TextRMSNorm(dimensions: config.hiddenSize, eps: config.rmsNormEps)
+        super.init()
+    }
+
+    func callAsFunction(_ hiddenStates: MLXArray, mask: MLXArray?) -> MLXArray {
+        var hidden = hiddenStates + self_attn(input_layernorm(hiddenStates), mask: mask)
+        hidden = hidden + mlp(post_attention_layernorm(hidden))
+        return hidden
+    }
+}
+
+private final class Qwen3Model: Module {
+    let config: Qwen3TextConfig
+    var memoryConfig: TextEncoderMemoryConfig = .aggressive
+
+    @ModuleInfo var embed_tokens: Embedding
+    var layers: [Qwen3DecoderLayer]
+    let norm: TextRMSNorm
+
+    init(config: Qwen3TextConfig) {
+        self.config = config
+        _embed_tokens = ModuleInfo(wrappedValue: Embedding(embeddingCount: config.vocabSize, dimensions: config.hiddenSize))
+        layers = (0..<config.numHiddenLayers).map { _ in Qwen3DecoderLayer(config: config) }
+        norm = TextRMSNorm(dimensions: config.hiddenSize, eps: config.rmsNormEps)
+        super.init()
+    }
+
+    func forwardWithHiddenStates(
+        _ inputIds: MLXArray,
+        layerIndices: [Int],
+        attentionMask: MLXArray
+    ) -> [Int: MLXArray] {
+        var hiddenStates = embed_tokens(inputIds)
+        eval(hiddenStates)
+
+        let mask = createCausalMask(seqLen: inputIds.shape[1], attentionMask: attentionMask)
+        eval(mask)
+
+        let layerSet = Set(layerIndices)
+        var extractedStates: [Int: MLXArray] = [:]
+
+        for (index, layer) in layers.enumerated() {
+            hiddenStates = layer(hiddenStates, mask: mask)
+            let layerIndex = index + 1
+            let shouldEval = layerSet.contains(layerIndex) || (index + 1) % memoryConfig.evalFrequency == 0
+
+            if shouldEval {
+                eval(hiddenStates)
+                if memoryConfig.clearCacheOnEval {
+                    MLX.Memory.clearCache()
+                }
+            }
+
+            if layerSet.contains(layerIndex) {
+                extractedStates[layerIndex] = hiddenStates
+            }
+        }
+
+        if layerSet.contains(layers.count) {
+            let normalizedStates = norm(hiddenStates)
+            eval(normalizedStates)
+            extractedStates[layers.count] = normalizedStates
+        }
+
+        MLX.Memory.clearCache()
+        return extractedStates
+    }
+
+    private func createCausalMask(seqLen: Int, attentionMask: MLXArray) -> MLXArray {
+        let rowIndices = MLXArray.arange(seqLen, dtype: .float32).expandedDimensions(axis: 1)
+        let colIndices = MLXArray.arange(seqLen, dtype: .float32).expandedDimensions(axis: 0)
+        let causalMask = MLX.where(
+            colIndices .<= rowIndices,
+            MLXArray(Float(0.0)),
+            MLXArray(-Float.infinity)
+        ).reshaped([1, 1, seqLen, seqLen])
+
+        let paddingMask = MLX.where(
+            attentionMask .== Int32(1),
+            MLXArray(Float(0.0)),
+            MLXArray(Float(-1e9))
+        ).reshaped([attentionMask.shape[0], 1, 1, attentionMask.shape[1]])
+
+        return causalMask + paddingMask
+    }
+}
+
+private final class Qwen3TextEncoderModel: Module {
+    var model: Qwen3Model
+
+    init(config: Qwen3TextConfig) {
+        model = Qwen3Model(config: config)
+        super.init()
+    }
+}
+
+public final class KleinTextEncoder: @unchecked Sendable {
+    private static let hiddenStateLayers = [9, 18, 27]
+
+    private let modelPath: URL
+    private var model: Qwen3TextEncoderModel?
+
+    public init(modelPath: URL) {
+        self.modelPath = modelPath
+    }
+
+    public func encode(inputIds: [Int], attentionMask: [Int]) throws -> MLXArray {
+        guard inputIds.count == attentionMask.count else {
+            throw Flux2Error.invalidConfiguration("inputIds and attentionMask must have the same length")
+        }
+
+        let model = try loadModel()
+        let ids = MLXArray(inputIds.map(Int32.init)).reshaped([1, inputIds.count])
+        let mask = MLXArray(attentionMask.map(Int32.init)).reshaped([1, attentionMask.count])
+        let extractedStates = model.model.forwardWithHiddenStates(
+            ids,
+            layerIndices: Self.hiddenStateLayers,
+            attentionMask: mask
+        )
+
+        let states = try Self.hiddenStateLayers.map { layerIndex -> MLXArray in
+            guard let state = extractedStates[layerIndex] else {
+                throw Flux2Error.invalidConfiguration("Qwen3 text encoder did not return hidden state layer \(layerIndex)")
+            }
+            return state
+        }
+
+        let embeddings = concatenated(states, axis: -1)
+        eval(embeddings)
+
+        // The text encoder is only needed to materialize prompt embeddings.
+        // Drop its 4B weights immediately so Flux image generation can load
+        // the transformer without keeping both models resident.
+        self.model = nil
+        MLX.Memory.clearCache()
+        return embeddings
+    }
+
+    private func loadModel() throws -> Qwen3TextEncoderModel {
+        if let model {
+            return model
+        }
+
+        let configURL = modelPath.appendingPathComponent("config.json")
+        let config = try Qwen3TextConfig.load(from: configURL)
+        let loaded = Qwen3TextEncoderModel(config: config)
+
+        let configData = try Data(contentsOf: configURL)
+        if let quantization = try? JSONDecoder().decode(Qwen3ModelConfigWithQuantization.self, from: configData).quantization {
+            quantize(model: loaded, groupSize: quantization.groupSize, bits: quantization.bits)
+        }
+
+        var weights: [String: MLXArray] = [:]
+        for file in try FileManager.default.contentsOfDirectory(atPath: modelPath.path).filter({ $0.hasSuffix(".safetensors") }).sorted() {
+            let tensors = try loadArrays(url: modelPath.appendingPathComponent(file))
+            for (key, value) in tensors {
+                weights[key] = value
+            }
+        }
+
+        guard !weights.isEmpty else {
+            throw Flux2Error.modelNotLoaded("Qwen3 text encoder weights not found")
+        }
+
+        try loaded.update(parameters: ModuleParameters.unflattened(weights), verify: .none)
+        eval(loaded)
+        MLX.Memory.clearCache()
+        model = loaded
+        return loaded
+    }
+}