diff --git a/Package.swift b/Package.swift index 424ceb9..626638c 100644 --- a/Package.swift +++ b/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"], diff --git a/Sources/FluidUse/EmbeddingGemma2/EmbeddingGemma2Manager.swift b/Sources/FluidUse/EmbeddingGemma2/EmbeddingGemma2Manager.swift index c245b1f..d311faf 100644 --- a/Sources/FluidUse/EmbeddingGemma2/EmbeddingGemma2Manager.swift +++ b/Sources/FluidUse/EmbeddingGemma2/EmbeddingGemma2Manager.swift @@ -216,13 +216,22 @@ public final class EmbeddingGemma2Manager: Sendable { /// Embedding of one audio window: ` <|audio> tokens ` 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: ` <|image> tokens `, `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.. static let audioEndId: Int32 = 258_883 // + static let imageStartId: Int32 = 255_999 // <|image> + static let imageEndId: Int32 = 258_882 // private var bosId: Int32 { tokenizer.bosId } private var eosId: Int32 { tokenizer.eosId } diff --git a/Sources/FluidUse/EmbeddingGemma2/EmbeddingGemma2ModelStore.swift b/Sources/FluidUse/EmbeddingGemma2/EmbeddingGemma2ModelStore.swift index 4bd23fb..25b118d 100644 --- a/Sources/FluidUse/EmbeddingGemma2/EmbeddingGemma2ModelStore.swift +++ b/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? diff --git a/Sources/FluidUse/EmbeddingGemma2/EmbeddingGemma2Vision.swift b/Sources/FluidUse/EmbeddingGemma2/EmbeddingGemma2Vision.swift new file mode 100644 index 0000000..23f9a9a --- /dev/null +++ b/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.. [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.. [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 + } +} diff --git a/Sources/ImageSearchCheck/main.swift b/Sources/ImageSearchCheck/main.swift new file mode 100644 index 0000000..df4d421 --- /dev/null +++ b/Sources/ImageSearchCheck/main.swift @@ -0,0 +1,72 @@ +import FluidUse +import Foundation + +/// Embeds images with EmbeddingGemma 2 and scores zero-shot Oxford Pets (photos cached by ImageSortDemo). +/// +/// swift run -c release ImageSearchCheck [--budget=70|140|280] [--per-breed=10] [--dump=vectors.jsonl] +@main +struct ImageSearchCheck { + static let breeds = [ + "abyssinian", "american bulldog", "american pit bull terrier", "basset hound", "beagle", "bengal", "birman", + "bombay", "boxer", "british shorthair", "chihuahua", "egyptian mau", "english cocker spaniel", + "english setter", "german shorthaired", "great pyrenees", "havanese", "japanese chin", "keeshond", + "leonberger", "maine coon", "miniature pinscher", "newfoundland", "persian", "pomeranian", "pug", "ragdoll", + "russian blue", "saint bernard", "samoyed", "scottish terrier", "shiba inu", "siamese", "sphynx", + "staffordshire bull terrier", "wheaten terrier", "yorkshire terrier", + ] + + static func main() async throws { + let arguments = CommandLine.arguments.dropFirst() + func value(_ name: String) -> String? { + arguments.first { $0.hasPrefix("--\(name)=") }.map { String($0.dropFirst(name.count + 3)) } + } + let budget = value("budget").flatMap(Int.init).flatMap(EmbeddingGemma2Vision.Budget.init) ?? .fast + let perBreed = value("per-breed").flatMap(Int.init) ?? 10 + let directory = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask)[0] + .appendingPathComponent("FluidUse/image-sort/oxford-pets-test") + let manifest = try JSONDecoder().decode( + [String: Int].self, from: Data(contentsOf: directory.appendingPathComponent("manifest.json"))) + // Same sample as the PyTorch check: test split (ids < 3,669), the first `perBreed` ids of each breed. + var byBreed: [Int: [Int]] = [:] + for (key, label) in manifest { if let id = Int(key), id < 3669 { byBreed[label, default: []].append(id) } } + let sample = (0.. Double { Double(DispatchTime.now().uptimeNanoseconds - start) / 1e9 } +} diff --git a/Sources/TileChallengeDemo/ChallengeModel.swift b/Sources/TileChallengeDemo/ChallengeModel.swift new file mode 100644 index 0000000..4e219ae --- /dev/null +++ b/Sources/TileChallengeDemo/ChallengeModel.swift @@ -0,0 +1,292 @@ +import CoreGraphics +import FluidUse +import Foundation +import SwiftUI + +/// Solves "select all images with …" photo grids with EmbeddingGemma 2: each tile is embedded (vision on the GPU, +/// text model on the Neural Engine) and ticked when the asked-for category beats all the others. Photos and labels are +/// 21 everyday categories from Caltech-256; the grids are built here, not taken from any real CAPTCHA. +@MainActor +final class ChallengeModel: ObservableObject { + enum Phase: Equatable { + case loading(String) + case ready + case running(String) + case done + case failed(String) + } + + struct Tile: Identifiable { + let id: Int + let image: CGImage + let category: Int + var picked: Bool? = nil + var score: Float = 0 + var isTarget: Bool + } + + struct Grid { + let prompt: String + var tiles: [Tile] + var verified = false + var solved: Bool { tiles.allSatisfy { $0.picked == $0.isTarget } } + } + + @Published private(set) var phase: Phase = .loading("Starting…") + @Published private(set) var grid: Grid? + @Published private(set) var history: [Bool] = [] + @Published private(set) var tilesDone = 0 + @Published private(set) var tilesCorrect = 0 + @Published private(set) var gridsSolved = 0 + @Published private(set) var busySeconds: Double = 0 + @Published private(set) var lastGridMilliseconds: Double = 0 + @Published private(set) var summary = "" + + /// Categories, and the lookalike groups the hard round draws its distractors from. + static let categories = [ + "traffic light", "fire hydrant", "bus", "bicycle", "motorcycle", "fire truck", "airplane", "helicopter", "boat", + "bridge", "lighthouse", "palm tree", "ladder", "umbrella", "horse", "zebra", "giraffe", "elephant", "camel", + "sunflower", "cactus", + ] + static let lookalikes: [[String]] = [ + ["bus", "bicycle", "motorcycle", "fire truck"], ["airplane", "helicopter", "boat"], + ["horse", "zebra", "giraffe", "elephant", "camel"], ["palm tree", "cactus", "sunflower"], + ["bridge", "lighthouse", "ladder", "traffic light", "fire hydrant", "umbrella"], + ] + static let turboGrids = max( + 1, CommandLine.arguments.first { $0.hasPrefix("--grids=") }.flatMap { Int($0.dropFirst(8)) } ?? 1000) + + private var text: EmbeddingGemma2Manager? + private var vision: EmbeddingGemma2Vision? + private var photos: [(file: URL, category: Int)] = [] + private var labelVectors: [[Float]] = [] // one per category + private var runTask: Task? + private var generator = SeededGenerator(seed: 7) + + var tilesPerSecond: Double { busySeconds > 0 ? Double(tilesDone) / busySeconds : 0 } + var accuracy: Double { tilesDone > 0 ? Double(tilesCorrect) / Double(tilesDone) : 0 } + var gridsDone: Int { history.count } + + func prepare() async { + do { + struct Item: Decodable { + let file: String + let label: String + } + let directory = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask)[0] + .appendingPathComponent("FluidUse/Datasets/caltech256-select") + let items = try JSONDecoder().decode( + [Item].self, from: Data(contentsOf: directory.appendingPathComponent("manifest.json"))) + photos = items.compactMap { item in + Self.categories.firstIndex(of: item.label).map { (directory.appendingPathComponent(item.file), $0) } + } + guard photos.count > 100 else { + throw EmbeddingGemma2Error.invalidAsset("No photos in \(directory.path)") + } + phase = .loading("Loading EmbeddingGemma 2…") + let loadStart = DispatchTime.now().uptimeNanoseconds + let text = try await EmbeddingGemma2Manager.loadDefault() + let vision = try await EmbeddingGemma2Vision.load(text: text) + for category in Self.categories { labelVectors.append(try await text.embed("a photo of a \(category).")) } + _ = try await vision.embed(image: try EmbeddingGemma2Vision.image(contentsOf: photos[0].file)) + self.text = text + self.vision = vision + DemoLog.model( + String( + format: "models loaded in %.1f s · vision on the GPU, text on the Neural Engine · %d photos", + Double(DispatchTime.now().uptimeNanoseconds - loadStart) / 1e9, photos.count)) + phase = .ready + if !CommandLine.arguments.contains("--manual") { start() } + } catch { + phase = .failed(error.localizedDescription) + DemoLog.line("failed: \(error.localizedDescription)", color: 196, bold: true) + } + } + + func start() { + guard phase == .ready || phase == .done else { return } + reset() + runTask = Task { await run() } + } + + func stop() { + runTask?.cancel() + runTask = nil + if case .running = phase { phase = .done } + } + + func reset() { + runTask?.cancel() + runTask = nil + grid = nil + history = [] + tilesDone = 0 + tilesCorrect = 0 + gridsSolved = 0 + busySeconds = 0 + summary = "" + generator = SeededGenerator(seed: 7) + if vision != nil { phase = .ready } + } + + private func run() async { + do { + DemoLog.event("▶ round 1: everyday things") + phase = .running("Round 1 · everyday things") + for _ in 0..<3 { try await showGrid(makeGrid(hard: false)) } + DemoLog.event("▶ round 2: lookalikes") + phase = .running("Round 2 · lookalikes (bus vs fire truck, horse vs zebra…)") + for _ in 0..<3 { try await showGrid(makeGrid(hard: true)) } + DemoLog.event("▶ turbo: \(Self.turboGrids) grids, no pauses") + phase = .running("Turbo · \(Self.turboGrids) grids, no pauses") + try await turbo() + summary = String( + format: "%d grids, %d tiles in %.1f s · %.0f tiles/s · %.1f%% of tiles right · %d/%d grids perfect", + gridsDone, tilesDone, busySeconds, tilesPerSecond, 100 * accuracy, gridsSolved, gridsDone) + DemoLog.event("■ " + summary) + phase = .done + } catch is CancellationError { + } catch { + phase = .failed(error.localizedDescription) + DemoLog.line("failed: \(error.localizedDescription)", color: 196, bold: true) + } + } + + typealias Spec = (prompt: String, target: Int, picks: [(URL, Int)]) + + /// A 3x3 grid with 2–5 photos of the target; distractors from any category, or (hard) from its lookalikes. + private func makeGrid(hard: Bool) -> Spec { + let positives = Int.random(in: 2...5, using: &generator) + let group = Self.lookalikes.randomElement(using: &generator)! + let pool = hard ? group.compactMap { Self.categories.firstIndex(of: $0) } : Array(Self.categories.indices) + let target = pool.randomElement(using: &generator)! + let yes = photos.filter { $0.category == target }.shuffled(using: &generator).prefix(positives) + let no = photos.filter { $0.category != target && pool.contains($0.category) }.shuffled(using: &generator) + .prefix(9 - yes.count) + let picks = (Array(yes) + Array(no)).shuffled(using: &generator).map { ($0.file, $0.category) } + let name = Self.categories[target] + return (("aeiou".contains(name.first!) ? "an " : "a ") + name, target, picks) + } + + /// Zero-shot decision: the asked-for category must beat every other category. + private func decide(_ vector: [Float], target: Int) -> (picked: Bool, score: Float) { + let scores = labelVectors.map { zip($0, vector).reduce(Float(0)) { $0 + $1.0 * $1.1 } } + let best = scores.indices.max { scores[$0] < scores[$1] } ?? 0 + return (best == target, scores[target]) + } + + private nonisolated static func decode(_ files: [URL]) throws -> [CGImage] { + try files.map { try EmbeddingGemma2Vision.image(contentsOf: $0) } + } + + /// A grid at a watchable pace: tiles appear, get ticked as each embedding lands, then Verify. + private func showGrid(_ spec: Spec) async throws { + guard let vision else { return } + let images = try await Task.detached { try Self.decode(spec.picks.map(\.0)) }.value + grid = Grid( + prompt: spec.prompt, + tiles: images.indices.map { + Tile(id: $0, image: images[$0], category: spec.picks[$0].1, isTarget: spec.picks[$0].1 == spec.target) + }) + try await Task.sleep(for: .milliseconds(450)) + let begin = DispatchTime.now().uptimeNanoseconds + let vectors = try await vision.embed(images: images, maxInFlight: 9) + let elapsed = Double(DispatchTime.now().uptimeNanoseconds - begin) / 1e9 + for index in vectors.indices { + let decision = decide(vectors[index], target: spec.target) + withAnimation(.easeOut(duration: 0.12)) { + grid?.tiles[index].picked = decision.picked + grid?.tiles[index].score = decision.score + } + try await Task.sleep(for: .milliseconds(70)) + } + record(elapsed: elapsed) + try await Task.sleep(for: .milliseconds(350)) + withAnimation(.easeOut(duration: 0.15)) { grid?.verified = true } + log() + try await Task.sleep(for: .milliseconds(700)) + } + + /// As fast as the chips allow: the next batch of grids is decoded and embedded while this one is shown. + private func turbo() async throws { + guard let vision else { return } + let batch = 4 + var remaining = Self.turboGrids + func prepareBatch() -> [Spec] { + let count = min(batch, remaining) + remaining -= count + return (0.. (images: [CGImage], vectors: [[Float]], seconds: Double) + { + let files = specs.flatMap { $0.picks.map(\.0) } + let begin = DispatchTime.now().uptimeNanoseconds + let images = try await Task.detached { try Self.decode(files) }.value + let vectors = try await vision.embed(images: images, maxInFlight: 12) + return (images, vectors, Double(DispatchTime.now().uptimeNanoseconds - begin) / 1e9) + } + var specs = prepareBatch() + var pending = Task { try await compute(specs) } + while !specs.isEmpty { + try Task.checkCancellation() + let (images, vectors, seconds) = try await pending.value + let current = specs + specs = prepareBatch() + let nextSpecs = specs + if !nextSpecs.isEmpty { pending = Task { try await compute(nextSpecs) } } + for (gridIndex, spec) in current.enumerated() { + var tiles: [Tile] = [] + for tileIndex in 0..<9 { + let flat = gridIndex * 9 + tileIndex + let decision = decide(vectors[flat], target: spec.target) + tiles.append( + Tile( + id: tileIndex, image: images[flat], category: spec.picks[tileIndex].1, + picked: decision.picked, + score: decision.score, isTarget: spec.picks[tileIndex].1 == spec.target)) + } + grid = Grid(prompt: spec.prompt, tiles: tiles, verified: true) + record(elapsed: seconds / Double(current.count)) + if gridsDone % 25 == 0 { log() } + try await Task.sleep(for: .milliseconds(16)) // one frame per grid, so every grid is drawn + } + } + } + + private func record(elapsed: Double) { + guard let grid else { return } + busySeconds += elapsed + lastGridMilliseconds = elapsed * 1000 + tilesDone += grid.tiles.count + tilesCorrect += grid.tiles.filter { $0.picked == $0.isTarget }.count + if grid.solved { gridsSolved += 1 } + history.append(grid.solved) + } + + private func log() { + guard let grid else { return } + DemoLog.line( + String( + format: "%@ grid %d “%@” %@ · %.0f ms · %.0f tiles/s · %.1f%% tiles right", grid.solved ? "✓" : "✗", + gridsDone, grid.prompt, grid.solved ? "solved" : "missed", lastGridMilliseconds, tilesPerSecond, + 100 * accuracy), color: grid.solved ? 120 : 203) + } +} + +/// Deterministic shuffles, so every run sees the same grids. +struct SeededGenerator: RandomNumberGenerator { + private var state: UInt64 + + init(seed: UInt64) { state = seed } + + mutating func next() -> UInt64 { + state &+= 0x9E37_79B9_7F4A_7C15 + var value = state + value = (value ^ (value >> 30)) &* 0xBF58_476D_1CE4_E5B9 + value = (value ^ (value >> 27)) &* 0x94D0_49BB_1331_11EB + return value ^ (value >> 31) + } +} diff --git a/Sources/TileChallengeDemo/ContentView.swift b/Sources/TileChallengeDemo/ContentView.swift new file mode 100644 index 0000000..1cc478f --- /dev/null +++ b/Sources/TileChallengeDemo/ContentView.swift @@ -0,0 +1,208 @@ +import SwiftUI + +struct ContentView: View { + @EnvironmentObject private var model: ChallengeModel + + var body: some View { + VStack(spacing: 0) { + Header() + Divider() + HStack(alignment: .top, spacing: 24) { + ChallengeCard() + VStack(alignment: .leading, spacing: 14) { + Stats() + HistoryStrip().frame(maxHeight: .infinity) + } + } + .padding(20) + } + .background(Color(nsColor: .windowBackgroundColor)) + } +} + +struct Header: View { + @EnvironmentObject private var model: ChallengeModel + + var body: some View { + HStack(spacing: 10) { + Text("Picture challenge").font(.title2.bold()) + Text(status).font(.callout).foregroundStyle(.secondary) + Spacer() + Button { + if case .running = model.phase { model.stop() } else { model.start() } + } label: { + if case .running = model.phase { + Label("Stop", systemImage: "pause.fill") + } else { + Label("Run", systemImage: "play.fill") + } + } + .keyboardShortcut(.space, modifiers: []) + .disabled(!canRun) + } + .padding(.horizontal, 20).padding(.vertical, 14) + } + + private var canRun: Bool { + switch model.phase { + case .ready, .done, .running: true + default: false + } + } + + private var status: String { + switch model.phase { + case .loading(let message): message + case .ready: "Ready · Space runs it" + case .running(let stage): stage + case .done: model.summary + case .failed(let message): message + } + } +} + +struct ChallengeCard: View { + @EnvironmentObject private var model: ChallengeModel + + var body: some View { + VStack(spacing: 0) { + VStack(alignment: .leading, spacing: 2) { + Text("Select all images with").font(.system(size: 15)) + Text(model.grid?.prompt ?? "…").font(.system(size: 26, weight: .bold)) + .contentTransition(.opacity) + } + .foregroundStyle(.white) + .frame(maxWidth: .infinity, alignment: .leading) + .padding(16) + .background(Color.indigo) + LazyVGrid(columns: Array(repeating: GridItem(.fixed(150), spacing: 4), count: 3), spacing: 4) { + ForEach(model.grid?.tiles ?? []) { tile in + TileView(tile: tile, verified: model.grid?.verified ?? false) + } + } + .frame(width: 458, height: 458) + .padding(4) + HStack { + if let grid = model.grid, grid.verified { + Label( + grid.solved ? "Solved" : "Missed", + systemImage: grid.solved ? "checkmark.seal.fill" : "xmark.seal.fill" + ) + .foregroundStyle(grid.solved ? .green : .red).font(.headline) + } + Spacer() + Text(String(format: "%.0f ms", model.lastGridMilliseconds)).font(.callout.monospacedDigit()) + .foregroundStyle(.secondary) + Text("VERIFY").font(.callout.bold()).foregroundStyle(.white) + .padding(.horizontal, 18).padding(.vertical, 8) + .background(RoundedRectangle(cornerRadius: 4).fill(Color.indigo)) + } + .padding(12) + } + .frame(width: 466) + .background(RoundedRectangle(cornerRadius: 6).fill(Color(nsColor: .controlBackgroundColor))) + .overlay(RoundedRectangle(cornerRadius: 6).stroke(Color.primary.opacity(0.15))) + } +} + +struct TileView: View { + let tile: ChallengeModel.Tile + let verified: Bool + + var body: some View { + Image(decorative: tile.image, scale: 1).resizable().scaledToFill() + .frame(width: 150, height: 150).clipped() + .scaleEffect(tile.picked == true ? 0.86 : 1) + .overlay(alignment: .topLeading) { + if tile.picked == true { + Image(systemName: "checkmark.circle.fill").font(.system(size: 26)) + .foregroundStyle(.white, Color.indigo).padding(6) + } + } + .overlay(alignment: .bottomTrailing) { + if tile.picked != nil { + Text(String(format: "%.2f", tile.score)).font(.caption2.monospacedDigit().bold()) + .padding(.horizontal, 5).padding(.vertical, 2) + .background(Capsule().fill(.black.opacity(0.55))).foregroundStyle(.white).padding(5) + } + } + .overlay { + if verified { + Rectangle().stroke(tile.picked == tile.isTarget ? Color.clear : Color.red, lineWidth: 5) + } + } + .background(Color.indigo.opacity(tile.picked == true ? 0.25 : 0)) + } +} + +struct Stats: View { + @EnvironmentObject private var model: ChallengeModel + + var body: some View { + LazyVGrid(columns: [GridItem(.flexible()), GridItem(.flexible())], spacing: 10) { + Tile(value: "\(model.gridsDone)", caption: "grids solved") + Tile(value: String(format: "%.0f", model.tilesPerSecond), caption: "photos / second") + Tile( + value: model.gridsDone == 0 ? "–" : String(format: "%.0f ms", model.lastGridMilliseconds), + caption: "per 9-photo grid") + Tile( + value: model.tilesDone == 0 ? "–" : String(format: "%.1f%%", 100 * model.accuracy), + caption: "photos judged right") + } + .frame(width: 440) + } +} + +/// Every grid so far, one square each (green solved, red missed), filling the space below the stats; squares +/// shrink as the run grows so nothing scrolls away. +struct HistoryStrip: View { + @EnvironmentObject private var model: ChallengeModel + + var body: some View { + VStack(alignment: .leading, spacing: 6) { + Text("\(model.gridsSolved) of \(model.gridsDone) grids perfect").font(.caption).foregroundStyle(.secondary) + Text( + "EmbeddingGemma 2 · photos on the GPU, labels on the Neural Engine · zero-shot · nothing leaves this Mac" + ) + .font(.caption2).foregroundStyle(.secondary) + Canvas { context, size in + let history = model.history + let count = max(history.count, 1) + // Largest square (gap = 1/3 of it, up to 12 pt) for which every grid fits in the space. + var cell: CGFloat = 12 + while cell > 2 { + let columns = max(Int(size.width / cell), 1) + if CGFloat((count + columns - 1) / columns) * cell <= size.height { break } + cell -= 0.5 + } + let columns = max(Int(size.width / cell), 1) + let square = cell * 0.75 + for (index, solved) in history.enumerated() { + let rect = CGRect( + x: CGFloat(index % columns) * cell, y: CGFloat(index / columns) * cell, width: square, + height: square) + context.fill( + Path(roundedRect: rect, cornerRadius: square * 0.2), with: .color(solved ? .green : .red)) + } + } + .frame(width: 440) + .frame(maxHeight: .infinity) + } + } +} + +struct Tile: View { + let value: String + let caption: String + + var body: some View { + VStack(alignment: .leading, spacing: 2) { + Text(value).font(.system(size: 30, weight: .semibold, design: .rounded)).monospacedDigit() + .contentTransition(.numericText()) + Text(caption).font(.caption).foregroundStyle(.secondary) + } + .frame(maxWidth: .infinity, alignment: .leading) + .padding(12) + .background(RoundedRectangle(cornerRadius: 10).fill(Color.primary.opacity(0.05))) + } +} diff --git a/Sources/TileChallengeDemo/DemoLog.swift b/Sources/TileChallengeDemo/DemoLog.swift new file mode 100644 index 0000000..7e5ddcd --- /dev/null +++ b/Sources/TileChallengeDemo/DemoLog.swift @@ -0,0 +1,29 @@ +import Foundation + +/// Colour-coded console lines for the side terminal (`demo.sh` tails stdout): model load, indexing, searches. +enum DemoLog { + /// 256-colour codes matching the collection colours. + private static let topicColors = [33, 208, 34, 205, 135, 37, 196, 61, 130, 121] + + static func line(_ text: String, color: Int? = nil, bold: Bool = false) { + let time = timeFormatter.string(from: Date()) + var styled = text + if let color { styled = "\u{1B}[38;5;\(color)m\(styled)\u{1B}[0m" } + if bold { styled = "\u{1B}[1m\(styled)\u{1B}[0m" } + print("\u{1B}[2m\(time)\u{1B}[0m \(styled)") + } + + static func colored(_ text: String, slot: Int) -> String { + "\u{1B}[38;5;\(topicColors[slot % topicColors.count])m\(text)\u{1B}[0m" + } + + static func event(_ text: String) { line(text, color: 220, bold: true) } + + static func model(_ text: String) { line("[model] " + text, color: 81) } + + private static let timeFormatter: DateFormatter = { + let formatter = DateFormatter() + formatter.dateFormat = "HH:mm:ss.SSS" + return formatter + }() +} diff --git a/Sources/TileChallengeDemo/README.md b/Sources/TileChallengeDemo/README.md new file mode 100644 index 0000000..8652301 --- /dev/null +++ b/Sources/TileChallengeDemo/README.md @@ -0,0 +1,17 @@ +# Picture challenge + +"Select all images with a …" photo grids, solved on the Mac by EmbeddingGemma 2 with no training: each photo is +embedded (vision encoder on the GPU, text model on the Neural Engine) and ticked when the asked-for category beats +all 21 (traffic light, fire hydrant, bus, bicycle, motorcycle, boat, bridge, palm tree, zebra, …). Photos and labels +are a Caltech-256 subset in `~/Library/Application Support/FluidUse/Datasets/caltech256-select` (`manifest.json` of +`{file, label}`); the grids are generated here, nothing comes from or is sent to any real CAPTCHA service. + +```bash +swift run -c release TileChallengeDemo # first run downloads ~0.9 GB of models +Sources/TileChallengeDemo/demo.sh # same, plus one terminal: macmon above the live log +``` + +It runs on its own: three mixed grids and three lookalike grids (bus vs fire truck, horse vs zebra…) at a watchable pace, then **turbo** +(`--grids=1000` by default): the next batch is decoded and embedded while the current grid is shown, and the counters +show photos per second, milliseconds per grid and accuracy against the true labels. Space stops or reruns it; +`--manual` waits for Space. diff --git a/Sources/TileChallengeDemo/TileChallengeDemoApp.swift b/Sources/TileChallengeDemo/TileChallengeDemoApp.swift new file mode 100644 index 0000000..54fec25 --- /dev/null +++ b/Sources/TileChallengeDemo/TileChallengeDemoApp.swift @@ -0,0 +1,24 @@ +import AppKit +import SwiftUI + +@main +struct TileChallengeDemoApp: App { + @StateObject private var model = ChallengeModel() + + init() { + setvbuf(stdout, nil, _IOLBF, 0) + // Bare SwiftPM executables start as background processes; make this one a regular windowed app. + NSApplication.shared.setActivationPolicy(.regular) + NSApplication.shared.activate(ignoringOtherApps: true) + } + + var body: some Scene { + WindowGroup("Picture challenge — on-device") { + ContentView() + .environmentObject(model) + .frame(minWidth: 980, minHeight: 680) + .task { await model.prepare() } + } + .defaultSize(width: 1120, height: 760) + } +} diff --git a/Sources/TileChallengeDemo/demo.sh b/Sources/TileChallengeDemo/demo.sh new file mode 100755 index 0000000..479da07 --- /dev/null +++ b/Sources/TileChallengeDemo/demo.sh @@ -0,0 +1,34 @@ +#!/bin/zsh +# Picture challenge demo: the app, plus one terminal (Ghostty if installed) with macmon (GPU / ANE / power; asitop as a +# fallback) on top and the live Neural Engine log below. +# Sources/TileChallengeDemo/demo.sh [app options, e.g. --autostart --auto-split --rate=200] +set -e +ROOT=$(cd "$(dirname "$0")/../.." && pwd) +LOG=${TMPDIR:-/tmp}/tile-challenge-demo.log +# macmon (brew install macmon) needs no sudo and works on M5; asitop 0.0.24 crashes there (KeyError E0-Cluster_active). +if command -v macmon >/dev/null; then + MONITOR="$(command -v macmon)" +else + MONITOR="sudo sh -c 'pkill -x powermetrics; exec $(command -v asitop || echo "$HOME/.local/bin/asitop")'" +fi + +swift build -c release --product TileChallengeDemo --package-path "$ROOT" +: > "$LOG" +# One tmux session, reused across launches: the monitor on top, the log below. +if ! tmux has-session -t tile-challenge 2>/dev/null; then + # The pane stays open if the monitor exits, so the reason is visible. + tmux new-session -d -s tile-challenge -x 160 -y 70 \ + "$MONITOR; echo; echo \"monitor exited (\$?) — press Enter to retry\"; read; $MONITOR" + tmux set-option -t tile-challenge remain-on-exit on + tmux split-window -v -t tile-challenge "tail -n 300 -F '$LOG'" + if [[ -d /Applications/Ghostty.app ]]; then + open -na Ghostty --args -e "$(command -v tmux)" attach -t tile-challenge + else + osascript -e 'tell application "Terminal" to do script "tmux attach -t tile-challenge"' \ + -e 'tell application "Terminal" to activate' >/dev/null + fi +fi +pkill -f "release/TileChallengeDemo" 2>/dev/null || true +cd "$ROOT" +"$ROOT/.build/release/TileChallengeDemo" "$@" > "$LOG" 2> "$LOG.stderr" & +echo "demo pid $! · log $LOG · runs on its own (Space stops / reruns)" diff --git a/Tests/FluidUseTests/EmbeddingGemma2Tests.swift b/Tests/FluidUseTests/EmbeddingGemma2Tests.swift index b0746d4..66b956d 100644 --- a/Tests/FluidUseTests/EmbeddingGemma2Tests.swift +++ b/Tests/FluidUseTests/EmbeddingGemma2Tests.swift @@ -109,3 +109,19 @@ final class EmbeddingGemma2AudioWindowTests: XCTestCase { XCTAssertEqual(EmbeddingGemma2Audio.validFrames(samples: 100), 0) } } + +/// Image resize targets, against Hugging Face `get_aspect_ratio_preserving_size`. +final class EmbeddingGemma2VisionSizeTests: XCTestCase { + func testTargetSizesMatchHuggingFace() { + let cases: [(width: Int, height: Int, budget: EmbeddingGemma2Vision.Budget, expected: (Int, Int))] = [ + (500, 334, .detailed, (960, 624)), (334, 500, .fast, (288, 480)), (1920, 1080, .balanced, (720, 384)), + (4000, 100, .fast, (2496, 48)), (64, 64, .detailed, (768, 768)), + ] + for item in cases { + let size = EmbeddingGemma2Vision.targetSize(width: item.width, height: item.height, budget: item.budget) + XCTAssertEqual(size.width, item.expected.0, "\(item.width)x\(item.height)") + XCTAssertEqual(size.height, item.expected.1, "\(item.width)x\(item.height)") + XCTAssertLessThanOrEqual((size.width / 16) * (size.height / 16), item.budget.patches) + } + } +}