EmbeddingGemma 2 images: ViT on Core ML, 79 images/s

Package.swift

@@ -54,6 +54,8 @@ let package = Package(
         .executableTarget(name: "BookmarkSortCheck", dependencies: ["BookmarkSort", "FluidUse"]),
         .executableTarget(name: "BookmarkSortDemo", dependencies: ["BookmarkSort"]),
         .executableTarget(name: "AudioSearchCheck", dependencies: ["FluidUse"]),
+        .executableTarget(name: "ImageSearchCheck", dependencies: ["FluidUse"]),
+        .executableTarget(name: "TileChallengeDemo", dependencies: ["FluidUse"], exclude: ["README.md", "demo.sh"]),
         .executableTarget(name: "AudioSearchDemo", dependencies: ["FluidUse"], exclude: ["README.md", "demo.sh"]),
         .executableTarget(
             name: "TopicSortDemo", dependencies: ["BookmarkSort", "FluidUse"], exclude: ["README.md", "demo.sh"],

Sources/FluidUse/EmbeddingGemma2/EmbeddingGemma2Manager.swift

@@ -216,13 +216,22 @@ public final class EmbeddingGemma2Manager: Sendable {
     /// Embedding of one audio window: `<bos> <|audio> tokens <audio|> <eos>` through the text model, where `tokens`
     /// are the first `count` rows of an `EmbeddingGemma2Audio` output ([1, 250, 512], fp16 or fp32).
     public func embed(audioTokens tokens: MLMultiArray, count: Int) async throws -> [Float] {
+        try await embed(media: tokens, count: count, start: Self.audioStartId, end: Self.audioEndId)
+    }
+
+    /// Embedding of one image: `<bos> <|image> tokens <image|> <eos>`, `tokens` from `EmbeddingGemma2Vision`.
+    public func embed(imageTokens tokens: MLMultiArray, count: Int) async throws -> [Float] {
+        try await embed(media: tokens, count: count, start: Self.imageStartId, end: Self.imageEndId)
+    }
+
+    private func embed(media tokens: MLMultiArray, count: Int, start: Int32, end: Int32) async throws -> [Float] {
         let total = count + 4
         guard let length = Self.lengths.first(where: { $0 >= total }) else {
-            throw EmbeddingGemma2Error.predictionFailed("\(count) audio tokens do not fit")
+            throw EmbeddingGemma2Error.predictionFailed("\(count) media tokens do not fit")
         }
         let tokenStride = tokens.strides[tokens.strides.count - 2].intValue
         return try await run(length: length, count: total) { destination in
-            writeEmbeddings([bosId, Self.audioStartId], to: destination)
+            writeEmbeddings([bosId, start], to: destination)
             let target = destination + 2 * Self.hiddenSize
             for row in 0..<count {
                 for column in 0..<Self.hiddenSize {
@@ -234,13 +243,15 @@ public final class EmbeddingGemma2Manager: Sendable {
                 }
             }
             writeEmbeddings(
-                [Self.audioEndId, eosId] + [Int32](repeating: tokenizer.padId, count: length - total),
+                [end, eosId] + [Int32](repeating: tokenizer.padId, count: length - total),
                 to: destination + (2 + count) * Self.hiddenSize)
         }
     }
 
     static let audioStartId: Int32 = 256_000  // <|audio>
     static let audioEndId: Int32 = 258_883  // <audio|>
+    static let imageStartId: Int32 = 255_999  // <|image>
+    static let imageEndId: Int32 = 258_882  // <image|>
     private var bosId: Int32 { tokenizer.bosId }
     private var eosId: Int32 { tokenizer.eosId }
 

Sources/FluidUse/EmbeddingGemma2/EmbeddingGemma2ModelStore.swift

@@ -31,6 +31,22 @@ public enum EmbeddingGemma2ModelStore {
             sha256: "41d7b552caabe376e377dcbe1c8c5f9879853bc4f384a161bb1f886f19e84b56"),
     ]
 
+    static let visionRevision = "f9d567d5f36c0a98773cdba26034b8185a353834"
+
+    private static let visionAssets = [
+        Asset(
+            path: "position_embeddings.f16", sha256: "17de9480acc23638c0634a97ed89e51ed664433256d7d6f109587dcd91c79124"),
+        Asset(
+            path: "EmbeddingGemma2Vision.mlpackage/Manifest.json",
+            sha256: "376afe7b51608637312cd15a0a89b1d7863f9779502aaf5f12fa6b14c4272d8d"),
+        Asset(
+            path: "EmbeddingGemma2Vision.mlpackage/Data/com.apple.CoreML/model.mlmodel",
+            sha256: "0085b5e04a13d000fe5969adeaa13d897f5a3cb7be60c76c8685c023747f89ef"),
+        Asset(
+            path: "EmbeddingGemma2Vision.mlpackage/Data/com.apple.CoreML/weights/weight.bin",
+            sha256: "5b5b6dd8d80aa438bf91ba143fbe90d56a2e25bfcf685a56b990aad71cf75232"),
+    ]
+
     private static let assets = [
         Asset(path: "config.json", sha256: "4ff891a28a5d55562b55ff3aa4ef362968d378629f9453f337cb70a9e1b45f61"),
         Asset(path: "tokenizer.json", sha256: "4d777ef5bdc1aa36227abdfb77c3e49e7b9c892d16e1b6bda41c393504828be4"),
@@ -62,6 +78,14 @@ public enum EmbeddingGemma2ModelStore {
             verifiedName: ".verified-audio-\(audioRevision)", cacheDirectory: cacheDirectory, progress: progress)
     }
 
+    /// Ensures the vision model (`EmbeddingGemma2Vision`) and its position tables exist, match their checksums, and
+    /// are compiled; returns the directory. Image tokens are embedded by the text model, so `ensure` is needed too.
+    public static func ensureVision(cacheDirectory: URL? = nil, progress: Progress? = nil) async throws -> URL {
+        try await ensure(
+            assets: visionAssets, revision: visionRevision, package: "EmbeddingGemma2Vision",
+            verifiedName: ".verified-vision-\(visionRevision)", cacheDirectory: cacheDirectory, progress: progress)
+    }
+
     private static func ensure(
         assets: [Asset], revision: String, package: String, verifiedName: String, cacheDirectory: URL?,
         progress: Progress?

Sources/FluidUse/EmbeddingGemma2/EmbeddingGemma2Vision.swift

@@ -0,0 +1,212 @@
+@preconcurrency import CoreML
+import CoreGraphics
+import Foundation
+import ImageIO
+
+/// EmbeddingGemma 2 image embeddings: the image is resized and cut into 16 px patches the way the Hugging Face
+/// `Gemma4ImageProcessor` does, encoded by `EmbeddingGemma2Vision` (ViT + 3x3 pooling, on the GPU, where it beats the
+/// Neural Engine at every size), and the pooled tokens are embedded by the text model, so images share the text space.
+public final class EmbeddingGemma2Vision: Sendable {
+    /// Soft tokens per image. More tokens keep more detail and cost more time (zero-shot Oxford Pets on 370 photos:
+    /// 87.0% / 88.4% / 89.5%; M5 Pro GPU: about 15 / 34 / 64+ ms per image).
+    public enum Budget: Int, CaseIterable, Sendable {
+        case fast = 70
+        case balanced = 140
+        case detailed = 280
+
+        var patches: Int { rawValue * 9 }
+    }
+
+    static let patchSize = 16
+    static let pooling = 3
+    static let hidden = 768
+    static let tableRows = 1024
+    public static let maxInFlight = 4
+
+    public let text: EmbeddingGemma2Manager
+    private let models: [Budget: MLModel]
+    private let positionTable: Data
+
+    init(text: EmbeddingGemma2Manager, models: [Budget: MLModel], positionTable: Data) {
+        self.text = text
+        self.models = models
+        self.positionTable = positionTable
+    }
+
+    /// Downloads (once, checksum-verified) and loads the vision model next to an already loaded text model.
+    public static func load(
+        text: EmbeddingGemma2Manager, computeUnits: MLComputeUnits = .cpuAndGPU,
+        progress: EmbeddingGemma2ModelStore.Progress? = nil
+    ) async throws -> EmbeddingGemma2Vision {
+        guard #available(macOS 15, iOS 18, *) else {
+            throw EmbeddingGemma2Error.unsupported("multifunction Core ML models need macOS 15 / iOS 18")
+        }
+        let directory: URL
+        if let path = ProcessInfo.processInfo.environment["EMBEDDINGGEMMA2_MODEL_DIR"], !path.isEmpty {
+            directory = URL(fileURLWithPath: path)
+        } else {
+            directory = try await EmbeddingGemma2ModelStore.ensureVision(progress: progress)
+        }
+        let compiled = directory.appendingPathComponent("EmbeddingGemma2Vision.mlmodelc")
+        let package = directory.appendingPathComponent("EmbeddingGemma2Vision.mlpackage")
+        let url: URL
+        if FileManager.default.fileExists(atPath: compiled.path) {
+            url = compiled
+        } else if FileManager.default.fileExists(atPath: package.path) {
+            url = try await MLModel.compileModel(at: package)
+        } else {
+            throw EmbeddingGemma2Error.invalidAsset("Missing EmbeddingGemma2Vision in \(directory.path)")
+        }
+        guard
+            let table = try? Data(
+                contentsOf: directory.appendingPathComponent("position_embeddings.f16"), options: .alwaysMapped),
+            table.count == 2 * tableRows * hidden * 2
+        else { throw EmbeddingGemma2Error.invalidAsset("Missing or malformed position_embeddings.f16") }
+        var models: [Budget: MLModel] = [:]
+        for budget in Budget.allCases {
+            let configuration = MLModelConfiguration()
+            configuration.computeUnits = computeUnits
+            configuration.functionName = "vision_\(budget.rawValue)"
+            models[budget] = try await MLModel.load(contentsOf: url, configuration: configuration)
+        }
+        return EmbeddingGemma2Vision(text: text, models: models, positionTable: table)
+    }
+
+    /// Decodes an image file (anything ImageIO reads).
+    public static func image(contentsOf url: URL) throws -> CGImage {
+        guard let source = CGImageSourceCreateWithURL(url as CFURL, nil),
+            let image = CGImageSourceCreateImageAtIndex(source, 0, nil)
+        else { throw EmbeddingGemma2Error.invalidAsset("Cannot decode \(url.lastPathComponent)") }
+        return image
+    }
+
+    /// Resized size in pixels: the largest aspect-preserving size whose sides are multiples of 48 px and which has at
+    /// most `budget.patches` 16 px patches (HF `get_aspect_ratio_preserving_size`).
+    static func targetSize(width: Int, height: Int, budget: Budget) -> (width: Int, height: Int) {
+        let side = patchSize * pooling
+        let factor = (Double(budget.patches * patchSize * patchSize) / Double(width * height)).squareRoot()
+        var targetWidth = Int((factor * Double(width) / Double(side)).rounded(.down)) * side
+        var targetHeight = Int((factor * Double(height) / Double(side)).rounded(.down)) * side
+        let maxSide = (budget.patches / (pooling * pooling)) * side
+        if targetHeight == 0 {
+            targetHeight = side
+            targetWidth = min((width / height) * side, maxSide)
+        } else if targetWidth == 0 {
+            targetWidth = side
+            targetHeight = min((height / width) * side, maxSide)
+        }
+        return (targetWidth, targetHeight)
+    }
+
+    /// Embeddings of many images, in order; up to `maxInFlight` run at once.
+    public func embed(
+        images: [CGImage], budget: Budget = .fast, maxInFlight: Int = EmbeddingGemma2Vision.maxInFlight
+    ) async throws -> [[Float]] {
+        try await withThrowingTaskGroup(of: (Int, [Float]).self) { group in
+            var result = [[Float]](repeating: [], count: images.count)
+            var next = 0
+            func addImage() {
+                guard next < images.count else { return }
+                let index = next
+                let image = images[index]
+                next += 1
+                group.addTask { (index, try await self.embed(image: image, budget: budget)) }
+            }
+            for _ in 0..<min(max(maxInFlight, 1), images.count) { addImage() }
+            for try await (index, vector) in group {
+                result[index] = vector
+                addImage()
+            }
+            return result
+        }
+    }
+
+    /// L2-normalized embedding of one image.
+    public func embed(image: CGImage, budget: Budget = .fast) async throws -> [Float] {
+        guard let model = models[budget] else { throw EmbeddingGemma2Error.predictionFailed("no vision model") }
+        let size = Self.targetSize(width: image.width, height: image.height, budget: budget)
+        let pixels = try Self.rgba(image, width: size.width, height: size.height)
+        let columns = size.width / Self.patchSize
+        let rows = size.height / Self.patchSize
+        let count = columns * rows
+        let total = budget.patches
+        let tokens = budget.rawValue
+        guard count <= total, columns <= Self.tableRows, rows <= Self.tableRows else {
+            throw EmbeddingGemma2Error.predictionFailed("image needs \(count) patches; the budget has \(total)")
+        }
+        let patches = try MLMultiArray(shape: [1, NSNumber(value: total), 768], dataType: .float16)
+        let positions = try MLMultiArray(shape: [1, NSNumber(value: total), 2], dataType: .float16)
+        let positionEmbeddings = try MLMultiArray(shape: [1, NSNumber(value: total), 768], dataType: .float16)
+        let valid = try MLMultiArray(shape: [1, NSNumber(value: total)], dataType: .float16)
+        let pool = try MLMultiArray(shape: [NSNumber(value: tokens), NSNumber(value: total)], dataType: .float16)
+        let patchPointer = patches.dataPointer.assumingMemoryBound(to: Float16.self)
+        let positionPointer = positions.dataPointer.assumingMemoryBound(to: Float16.self)
+        let embeddingPointer = positionEmbeddings.dataPointer.assumingMemoryBound(to: Float16.self)
+        let validPointer = valid.dataPointer.assumingMemoryBound(to: Float16.self)
+        let poolPointer = pool.dataPointer.assumingMemoryBound(to: Float16.self)
+        patchPointer.update(repeating: 0, count: total * 768)
+        embeddingPointer.update(repeating: 0, count: total * 768)
+        positionPointer.update(repeating: -1, count: total * 2)
+        validPointer.update(repeating: 0, count: total)
+        poolPointer.update(repeating: 0, count: tokens * total)
+        positionTable.withUnsafeBytes { raw in
+            let table = raw.bindMemory(to: Float16.self)
+            for row in 0..<rows {
+                for column in 0..<columns {
+                    let patch = row * columns + column
+                    // Patch layout: [16 rows][16 columns][RGB], values in [0, 1] (the model maps them to [-1, 1]).
+                    let base = patch * 768
+                    for y in 0..<Self.patchSize {
+                        for x in 0..<Self.patchSize {
+                            let pixel = ((row * Self.patchSize + y) * size.width + column * Self.patchSize + x) * 4
+                            let offset = base + (y * Self.patchSize + x) * 3
+                            for channel in 0..<3 {
+                                patchPointer[offset + channel] = Float16(Float(pixels[pixel + channel]) / 255)
+                            }
+                        }
+                    }
+                    positionPointer[patch * 2] = Float16(column)
+                    positionPointer[patch * 2 + 1] = Float16(row)
+                    validPointer[patch] = 1
+                    let xRow = column * Self.hidden
+                    let yRow = (Self.tableRows + row) * Self.hidden
+                    for channel in 0..<Self.hidden {
+                        embeddingPointer[base + channel] = Float16(
+                            Float(table[xRow + channel]) + Float(table[yRow + channel]))
+                    }
+                    // 3x3 average pooling over the patch grid, groups numbered row-major.
+                    let group = (column / Self.pooling) + (columns / Self.pooling) * (row / Self.pooling)
+                    poolPointer[group * total + patch] = Float16(1.0 / 9.0)
+                }
+            }
+        }
+        let input = try MLDictionaryFeatureProvider(dictionary: [
+            "patches": MLFeatureValue(multiArray: patches), "positions": MLFeatureValue(multiArray: positions),
+            "position_embeddings": MLFeatureValue(multiArray: positionEmbeddings),
+            "valid": MLFeatureValue(multiArray: valid), "pool": MLFeatureValue(multiArray: pool),
+        ])
+        let output = try await model.prediction(from: input)
+        guard let imageTokens = output.featureValue(for: "image_tokens")?.multiArrayValue else {
+            throw EmbeddingGemma2Error.predictionFailed("missing image_tokens output")
+        }
+        return try await text.embed(imageTokens: imageTokens, count: count / (Self.pooling * Self.pooling))
+    }
+
+    /// RGBA8 pixels of `image` drawn at `width` x `height` with high-quality interpolation.
+    static func rgba(_ image: CGImage, width: Int, height: Int) throws -> [UInt8] {
+        var pixels = [UInt8](repeating: 0, count: width * height * 4)
+        let drawn = pixels.withUnsafeMutableBytes { buffer -> Bool in
+            guard
+                let context = CGContext(
+                    data: buffer.baseAddress, width: width, height: height, bitsPerComponent: 8, bytesPerRow: width * 4,
+                    space: CGColorSpace(name: CGColorSpace.sRGB)!,
+                    bitmapInfo: CGImageAlphaInfo.noneSkipLast.rawValue)
+            else { return false }
+            context.interpolationQuality = .high
+            context.draw(image, in: CGRect(x: 0, y: 0, width: width, height: height))
+            return true
+        }
+        guard drawn else { throw EmbeddingGemma2Error.predictionFailed("cannot draw the image") }
+        return pixels
+    }
+}