diff --git a/Package.swift b/Package.swift index 90b883c..13b14ea 100644 --- a/Package.swift +++ b/Package.swift @@ -67,8 +67,13 @@ let package = Package( .executableTarget(name: "ImageSortDemo", dependencies: ["ImageSort"], exclude: ["README.md"]), .executableTarget(name: "KevCheck", dependencies: ["FluidUse"]), .executableTarget(name: "ClefVisionCheck", dependencies: ["FluidUse"]), - .executableTarget(name: "ClefFlashDemo", dependencies: ["FluidUse"], exclude: ["README.md", "demo.sh", "bars.py"]), + .executableTarget( + name: "ClefFlashDemo", dependencies: ["FluidUse"], exclude: ["README.md", "demo.sh", "bars.py"]), .executableTarget(name: "ClefTextCheck", dependencies: ["FluidUse"]), + .executableTarget(name: "JevCheck", dependencies: ["FluidUse"]), + .executableTarget( + name: "JevMarioDemo", dependencies: ["FluidUse"], + exclude: ["README.md", "demo.sh", "nes_bridge.py", "retro_bridge.py"]), .executableTarget(name: "ClefCompareDemo", dependencies: ["FluidUse"], exclude: ["README.md", "demo.sh"]), .executableTarget(name: "InternDecisionCheck", dependencies: ["FluidUse"]), .executableTarget(name: "Decision2Check", dependencies: ["FluidUse"]), @@ -79,7 +84,8 @@ let package = Package( .executableTarget(name: "GuardrailDemo", dependencies: ["FluidUse"], exclude: ["README.md"]), .executableTarget(name: "FrontDoorDemo", dependencies: ["FluidUse"], exclude: ["README.md"]), .executableTarget(name: "ShortReplyCheck", dependencies: ["FluidUse"]), - .executableTarget(name: "ShortReplyDemo", dependencies: ["FluidUse"], exclude: ["README.md", "demo.sh", "mock-feed"]), + .executableTarget( + name: "ShortReplyDemo", dependencies: ["FluidUse"], exclude: ["README.md", "demo.sh", "mock-feed"]), .executableTarget(name: "GLiClassServe", dependencies: ["FluidUse"]), .executableTarget( name: "KevGuessWhoDemo", dependencies: ["FluidUse", "SortAnything"], exclude: ["README.md", "demo.sh"]), diff --git a/Sources/ClefFlashDemo/bars.py b/Sources/ClefFlashDemo/bars.py index 4d1969e..952b8dc 100755 --- a/Sources/ClefFlashDemo/bars.py +++ b/Sources/ClefFlashDemo/bars.py @@ -1,16 +1,23 @@ #!/usr/bin/env python3 -"""Live bar charts for the demo's top pane: CPU / GPU utilisation and GPU power from `macmon pipe` (no sudo). +"""Live bar charts for the demo's top pane: CPU / GPU utilisation and GPU power from `macmon pipe` (no sudo); with +--ane also Neural Engine power from `sudo powermetrics` (asks for the password once). On M5 Pro / macOS 27, macmon's CPU and ANE power read 0 W and powermetrics' GPU power reads milliwatts while the GPU is busy, so this shows the counters that do track load: P-/E-cluster usage, GPU usage and GPU watts. """ +from __future__ import annotations + import json +import plistlib import shutil import subprocess import sys +import threading +import time BAR = "█" -COLORS = {"P-CPU": "\033[38;5;75m", "E-CPU": "\033[38;5;111m", "GPU": "\033[38;5;214m", "GPU W": "\033[38;5;203m"} +COLORS = {"P-CPU": "\033[38;5;75m", "E-CPU": "\033[38;5;111m", "GPU": "\033[38;5;214m", "GPU W": "\033[38;5;203m", + "ANE W": "\033[38;5;120m", "ANE": "\033[38;5;84m"} RESET, DIM, BOLD = "\033[0m", "\033[2m", "\033[1m" @@ -21,7 +28,93 @@ def bar(label: str, fraction: float, text: str, width: int) -> str: f"{text}") +def find_ane_watts(node) -> float | None: + """ANE power from a powermetrics plist sample (key names vary by macOS / chip: ane_power in mW, ane_energy ...).""" + if isinstance(node, dict): + for key, value in node.items(): + k = key.lower() + if "ane" in k and "power" in k and isinstance(value, (int, float)): + return float(value) / 1000 # mW + for value in node.values(): + found = find_ane_watts(value) + if found is not None: + return found + elif isinstance(node, list): + for value in node: + found = find_ane_watts(value) + if found is not None: + return found + return None + + +class AneMeter: + """`sudo powermetrics --samplers ane_power` (macmon's ANE power reads 0 W on macOS 27). Latest watts + peak.""" + + def __init__(self): + self.watts, self.peak, self.error = 0.0, 0.0, None + self.process = subprocess.Popen( + ["sudo", "powermetrics", "--samplers", "ane_power", "-i", "500", "-f", "plist"], stdout=subprocess.PIPE) + threading.Thread(target=self.read, daemon=True).start() + + def read(self): + buffer = b"" + for chunk in iter(lambda: self.process.stdout.read(4096), b""): + buffer += chunk + while b"\0" in buffer: + sample, buffer = buffer.split(b"\0", 1) + try: + watts = find_ane_watts(plistlib.loads(sample.strip())) + except Exception as e: # noqa: BLE001 + self.error = str(e)[:60] + continue + if watts is None: + self.error = "no ANE power key in powermetrics output" + continue + self.watts, self.peak = watts, max(self.peak, watts) + + +class AneBusy: + """Neural Engine busy time from the demo's own decision log (`... 0.55 ms` per line): the model's ops all run on + the ANE (MLComputePlan), so model milliseconds over the last second = ANE busy %. No sudo, no hardware counter.""" + + def __init__(self, path: str, window: float = 1.0): + self.window, self.events, self.placement = window, [], "" + try: # the placement line is printed once at load, possibly before this pane started + self.placement = next((l.split(",", 1)[-1].strip() for l in open(path) if "Neural Engine" in l), "") + except OSError: + pass + self.process = subprocess.Popen(["tail", "-n", "0", "-F", path], stdout=subprocess.PIPE, text=True, + stderr=subprocess.DEVNULL) + threading.Thread(target=self.read, daemon=True).start() + + def read(self): + for line in self.process.stdout: + parts = line.split() + if "Neural Engine" in line and "ops" in line: + self.placement = line.split(",", 1)[-1].strip() + if len(parts) >= 2 and parts[-1] == "ms": + try: + self.events.append((time.time(), float(parts[-2]))) + except ValueError: + pass + + def stats(self) -> tuple[float, float, float]: + """(busy fraction, median ms per decision, decisions per second) over the window.""" + now = time.time() + self.events = [e for e in self.events if now - e[0] <= max(self.window, 5.0)] + recent = [ms for t, ms in self.events if now - t <= self.window] + busy = sum(recent) / (self.window * 1000) + ms = sorted(recent)[len(recent) // 2] if recent else 0.0 + return busy, ms, len(recent) / self.window + + def main() -> None: + busy = AneBusy(sys.argv[sys.argv.index("--log") + 1]) if "--log" in sys.argv else None + ane = None + if "--ane" in sys.argv: + print("ANE power needs `sudo powermetrics`:", flush=True) + subprocess.run(["sudo", "-v"]) # password prompt here, before the screen starts redrawing + ane = AneMeter() process = subprocess.Popen(["macmon", "pipe", "-i", "500"], stdout=subprocess.PIPE, text=True) gpu_watts_max = 30.0 sys.stdout.write("\033[?25l") # hide cursor @@ -31,7 +124,8 @@ def main() -> None: d = json.loads(line) except json.JSONDecodeError: continue - width = max(10, shutil.get_terminal_size().columns - 30) + # label (9) + space + bar + space + the widest right-hand text, so nothing clips at any window width + width = max(10, shutil.get_terminal_size().columns - 9 - 2 - (24 if busy else 15) - 2) p_mhz, p_use = d["pcpu_usage"] e_mhz, e_use = d["ecpu_usage"] g_mhz, g_use = d["gpu_usage"] @@ -46,6 +140,21 @@ def main() -> None: bar("GPU", g_use, f"{g_use * 100:5.1f}% {g_mhz:4d} MHz", width), bar("GPU W", g_watts / gpu_watts_max, f"{g_watts:5.1f} W", width), ] + if busy is not None: + fraction, ms, rate = busy.stats() + # the bar is drawn at least one cell while the model runs: busy time is a fraction of a percent + cells = max(1 if rate else 0, int(round(fraction * width))) + rows.append(f" {BOLD}{'ANE':<6}{RESET} {COLORS['ANE']}{BAR * cells}{RESET}{DIM}{'·' * (width - cells)}{RESET} " + f"{fraction * 100:5.2f}% {ms:4.2f} ms × {rate:.0f}/s") + if busy.placement: + ops = busy.placement.split("·")[-1].strip() + rows.append(f" {DIM}ANE % = model time per second · {ops} on the Neural Engine{RESET}") + if ane is not None: + # no utilisation counter exists for the ANE; power vs. a 2 W scale (or the session peak) stands in + scale = max(2.0, ane.peak) + rows.append(bar("ANE W", ane.watts / scale, f"{ane.watts * 1000:5.0f} mW", width)) + if ane.error: + rows.append(f" {DIM}ANE: {ane.error}{RESET}") if ram: rows += ["", f" {DIM}RAM {ram.get('ram_usage', 0) / 2**30:.1f} / {ram.get('ram_total', 0) / 2**30:.0f} GB" f" swap {ram.get('swap_usage', 0) / 2**30:.1f} GB{RESET}"] diff --git a/Sources/FluidUse/Jev/JevManager.swift b/Sources/FluidUse/Jev/JevManager.swift new file mode 100644 index 0000000..dcdf228 --- /dev/null +++ b/Sources/FluidUse/Jev/JevManager.swift @@ -0,0 +1,312 @@ +import Accelerate +import CoreGraphics +import CoreML +import Foundation + +/// autotrust/JEV-9B with vision, System 1, on Core ML: typed decisions (yes/no, 0-5 score, choice) over text and +/// images, one calibrated probability per option from a single prefill. +/// +/// Packages: the Qwen3.5-9B vision tower (fp32, one per patch bucket), the decoder with JEV's System 1 LoRA merged in +/// (8 layer-group parts, 8-bit, one function per sequence bucket, chained on the GPU), and on the host the adapted +/// `lm_head` rows at the decision-token ids plus JEV's per-slot bias and per-kind temperature. The prompt is +/// JEV's `serve_decide.py` System 1 template. An actor: Core ML's synchronous `prediction` is not thread-safe. +@available(macOS 15.0, iOS 18.0, *) +public actor JevManager { + public nonisolated let bucket: Int + let config: JevConfig + let tokenizer: QwenBPETokenizer + private let parts: [MLModel] + private let embeddings: Data // fp16 [vocab, hidden] + private let headRows: [Float] // fp32 [headCount, hidden] + private let posEmbed: Data // fp32 vision position table + private let vision: [(patches: Int, model: MLModel)] // ascending patch buckets + + /// Loads one decoder sequence bucket (prompts up to `bucket` tokens) and the vision packages for `visionPatches` + /// (default: every bucket in the bundle). Packages are compiled once into `compiled/`. + public static func load( + from directory: URL, bucket: Int = 256, visionPatches: [Int]? = nil, computeUnits: MLComputeUnits = .cpuAndGPU + ) async throws -> JevManager { + func json(_ name: String) throws -> [String: Any] { + let object = try JSONSerialization.jsonObject( + with: Data(contentsOf: directory.appendingPathComponent(name))) + guard let dict = object as? [String: Any] else { throw JevError.invalidAsset(name) } + return dict + } + let config = try JevConfig(manifest: try json("config.json"), head: try json("head.json")) + guard config.buckets.contains(bucket) else { throw JevError.invalidInput("no \(bucket)-token bucket") } + let tokenizer = try QwenBPETokenizer(tokenizerJsonURL: directory.appendingPathComponent("tokenizer.json")) + let embeddings = try Data( + contentsOf: directory.appendingPathComponent("embeddings.f16"), options: .alwaysMapped) + guard embeddings.count == config.vocab * config.hidden * 2 else { + throw JevError.invalidAsset("embeddings.f16 size") + } + let rowsData = try Data(contentsOf: directory.appendingPathComponent(config.headRows)) + guard rowsData.count == config.headCount * config.hidden * 4 else { + throw JevError.invalidAsset("head rows size") + } + let headRows = rowsData.withUnsafeBytes { Array($0.bindMemory(to: Float.self)) } + let posEmbed = try Data(contentsOf: directory.appendingPathComponent(config.posEmbed)) + let v = config.clef.vision + guard posEmbed.count == v.depthPositions * v.hidden * 4 else { throw JevError.invalidAsset("pos_embed size") } + var parts: [MLModel] = [] + for path in config.parts { + parts.append(try await compiled(path, in: directory, units: computeUnits, function: "L\(bucket)")) + } + var vision: [(Int, MLModel)] = [] + for package in config.visionPackages where visionPatches?.contains(package.patches) ?? true { + vision.append( + (package.patches, try await compiled(package.file, in: directory, units: computeUnits, function: nil))) + } + return JevManager( + bucket: bucket, config: config, tokenizer: tokenizer, parts: parts, embeddings: embeddings, + headRows: headRows, posEmbed: posEmbed, vision: vision) + } + + static func compiled( + _ path: String, in directory: URL, units: MLComputeUnits, function: String? + ) async throws + -> MLModel + { + let compiledDir = directory.appendingPathComponent("compiled") + try FileManager.default.createDirectory(at: compiledDir, withIntermediateDirectories: true) + let source = directory.appendingPathComponent(path) + let target = compiledDir.appendingPathComponent(source.deletingPathExtension().lastPathComponent + ".mlmodelc") + if !FileManager.default.fileExists(atPath: target.path) { + let temporary = try await MLModel.compileModel(at: source) + try FileManager.default.moveItem(at: temporary, to: target) + } + let configuration = MLModelConfiguration() + configuration.computeUnits = units + configuration.functionName = function + return try await MLModel.load(contentsOf: target, configuration: configuration) + } + + init( + bucket: Int, config: JevConfig, tokenizer: QwenBPETokenizer, parts: [MLModel], embeddings: Data, + headRows: [Float], posEmbed: Data, vision: [(Int, MLModel)] + ) { + self.bucket = bucket + self.config = config + self.tokenizer = tokenizer + self.parts = parts + self.embeddings = embeddings + self.headRows = headRows + self.posEmbed = posEmbed + self.vision = vision.sorted { $0.0 < $1.0 }.map { (patches: $0.0, model: $0.1) } + } + + /// The smallest loaded vision bucket that holds `patches`. + private func visionModel(for patches: Int) throws -> (MLModel, Int) { + guard let package = vision.first(where: { $0.patches >= patches }) else { + throw JevError.invalidInput( + "image has \(patches) patches; the largest loaded vision bucket holds \(vision.last?.patches ?? 0)") + } + return (package.model, package.patches) + } + + // MARK: prompt + + struct Encoded { + let inputIDs: [Int] + let positions: [[Int]] + let imageRanges: [Range] + let optionCount: Int + } + + /// `serve_decide.s1_pass` text with each image expanded to its merged-token count. Adjacent text parts are joined + /// before tokenizing and images are spliced as ids, which matches Hugging Face splitting on added tokens. + func encode( + state: [JevStatePart], imageTokens: [Int], grids: [(h: Int, w: Int)], question: JevQuestion + ) throws + -> Encoded + { + let lines: [String] + switch question.kind { + case .noul: lines = ["false", "true"] + case .score: lines = (0...5).map(String.init) + case .choice(let options): + guard (2...config.labels.count).contains(options.count) else { + throw JevError.invalidInput("choice needs 2-\(config.labels.count) options, got \(options.count)") + } + lines = options.enumerated().map { "\(config.labels[$0.offset])) \($0.element)" } + } + var ids: [Int] = [] + var ranges: [Range] = [] + var pending = "[kind] \(question.kind.name)\n[state] " + var image = 0 + for part in state { + switch part { + case .text(let text): pending += text + case .image: + ids += try tokenizer.encode(pending) + pending = "" + ids.append(config.visionStartID) + let start = ids.count + ids += [Int](repeating: config.imageTokenID, count: imageTokens[image]) + ranges.append(start.. (ids: [Int], positions: [[Int]]) { + var counts: [Int] = [] + var grids: [(h: Int, w: Int)] = [] + for part in state { + guard case .image(let image) = part else { continue } + let patches = try ClefImagePreprocessor.patches(from: image, config: config.clef) + counts.append(patches.mergedTokens) + grids.append((patches.gridH, patches.gridW)) + } + let encoded = try encode(state: state, imageTokens: counts, grids: grids, question: question) + return (encoded.inputIDs, encoded.positions) + } + + // MARK: decide + + /// Answer every question about the same state. Images go through the vision tower once; each question is one + /// decoder pass. Probabilities follow the option order of each question (`["false", "true"]` for noul). + public func decide(state: [JevStatePart], questions: [JevQuestion]) throws -> JevResult { + let started = Date() + // 1. vision tokens per image + var visionTokens: [[Float]] = [] + var counts: [Int] = [] + var grids: [(h: Int, w: Int)] = [] + let visionStarted = Date() + for part in state { + guard case .image(let image) = part else { continue } + let patches = try ClefImagePreprocessor.patches(from: image, config: config.clef) + let (model, bucketPatches) = try visionModel(for: patches.count) + let clef = config.clef(patches: bucketPatches) + let inputs = try posEmbed.withUnsafeBytes { + try ClefVisionHost.visionInputs(patches, posEmbed: $0.bindMemory(to: Float.self), config: clef) + } + guard let out = try model.prediction(from: inputs.features).featureValue(for: "tokens")?.multiArrayValue + else { throw JevError.invalidAsset("vision tower returned no tokens") } + visionTokens.append(Self.rows(out, count: inputs.mergedTokens, width: config.hidden)) + counts.append(inputs.mergedTokens) + grids.append((patches.gridH, patches.gridW)) + } + let visionMs = Date().timeIntervalSince(visionStarted) * 1000 + // 2. one decoder pass per question + let (cos, sin) = try ropeArrays() // refilled per question + var decisions: [JevDecision] = [] + var tokens: [Int] = [] + var decoderMs = 0.0 + for question in questions { + let encoded = try encode(state: state, imageTokens: counts, grids: grids, question: question) + try fillRope(cos: cos, sin: sin, positions: encoded.positions) + var hidden = try hiddenRow(encoded, visionTokens: visionTokens) + let decoderStarted = Date() + for part in parts { + let output = try part.prediction( + from: MLDictionaryFeatureProvider(dictionary: ["hidden": hidden, "cos": cos, "sin": sin])) + guard let next = output.featureValue(for: "out")?.multiArrayValue else { + throw JevError.invalidAsset("decoder part returned no output") + } + hidden = next + } + decoderMs += Date().timeIntervalSince(decoderStarted) * 1000 + decisions.append(try decision(hidden, row: encoded.inputIDs.count - 1, kind: question.kind)) + tokens.append(encoded.inputIDs.count) + } + return JevResult( + decisions: decisions, inputTokens: tokens, visionMilliseconds: visionMs, decoderMilliseconds: decoderMs, + totalMilliseconds: Date().timeIntervalSince(started) * 1000) + } + + /// Head: adapted lm_head rows · last real row, + bias, / temperature, softmax over the question's options. + func decision(_ hidden: MLMultiArray, row: Int, kind: JevKind) throws -> JevDecision { + let D = config.hidden + let rowStride = hidden.strides[1].intValue + let colStride = hidden.strides[2].intValue + let src = hidden.dataPointer.assumingMemoryBound(to: Float16.self) + var last = [Float](repeating: 0, count: D) + for d in 0.. MLMultiArray { + let L = bucket + let D = config.hidden + let padID = config.padID + let ids = encoded.inputIDs + let array = try MLMultiArray(shape: [1, L as NSNumber, D as NSNumber], dataType: .float16) + embeddings.withUnsafeBytes { raw in + let table = raw.bindMemory(to: Float16.self).baseAddress! + let out = array.dataPointer.assumingMemoryBound(to: Float16.self) + for position in 0.. (MLMultiArray, MLMultiArray) { + let shape = [bucket as NSNumber, config.rotaryDim as NSNumber] + return ( + try MLMultiArray(shape: shape, dataType: .float16), try MLMultiArray(shape: shape, dataType: .float16) + ) + } + + /// Interleaved M-RoPE tables (fp16 [L, R]); padding continues the position count after the prompt. + func fillRope(cos: MLMultiArray, sin: MLMultiArray, positions: [[Int]]) throws { + let (c32, s32) = try ClefVisionHost.ropeTables(positions: positions, length: bucket, config: config.clef) + let count = bucket * config.rotaryDim + let c = cos.dataPointer.assumingMemoryBound(to: Float16.self) + let s = sin.dataPointer.assumingMemoryBound(to: Float16.self) + let cs = c32.dataPointer.assumingMemoryBound(to: Float.self) + let ss = s32.dataPointer.assumingMemoryBound(to: Float.self) + for i in 0.. [Float] { + let rowStride = array.strides[0].intValue + let colStride = array.strides[1].intValue + let src = array.dataPointer.assumingMemoryBound(to: Float.self) + var out = [Float](repeating: 0, count: count * width) + for r in 0.. Void + + struct Asset { + let path: String + let sha256: String + } + + public static let repository = "FluidInference/jev-9b-coreml" + public static let revision = "a718625da54446f09cd63d8c3962bfe53d81c4a9" + static let assets: [Asset] = [ + Asset(path: "config.json", sha256: "bccfe3e196fde9da7357fc3ef8f6714999ea9e60b2b7e94248603e123e201a69"), + Asset(path: "embeddings.f16", sha256: "7d87bcd3dbacdc2ab33fff49adaa4c4f35691d829c8aae75670bcf8792a10242"), + Asset(path: "head_rows.f32", sha256: "214244c905063010fb2473bca2ecb44645119e056cdf10a284e0f70ad3d9b21f"), + Asset(path: "head.json", sha256: "a06764e3018281e46ad48f107772f40880dee9e5fbd363582e2df17290a24835"), + Asset(path: "LICENSE", sha256: "bbedc3fda3305820b977265f01b8619d87570a6739de3a5582c3464840f1e57a"), + Asset( + path: "part00.mlpackage/Data/com.apple.CoreML/model.mlmodel", + sha256: "0bcab81eeb64a8cb38f30ad5999f91da64291a7973431fabeb595978ddd30634"), + Asset( + path: "part00.mlpackage/Data/com.apple.CoreML/weights/weight.bin", + sha256: "d6afcc3b6c7a9747247851778792a5c13f6b244d53eb1cff717a3b7881b271f2"), + Asset( + path: "part00.mlpackage/Manifest.json", + sha256: "c27ee94412bcd5514f120c9848f0eed75c51322714cc4912e22bf23a8ada513a"), + Asset( + path: "part01.mlpackage/Data/com.apple.CoreML/model.mlmodel", + sha256: "0bcab81eeb64a8cb38f30ad5999f91da64291a7973431fabeb595978ddd30634"), + Asset( + path: "part01.mlpackage/Data/com.apple.CoreML/weights/weight.bin", + sha256: "02c09420f6dc3eb1ac93c2adb3c39adede616beabfad0a19e6ef27d3c1117202"), + Asset( + path: "part01.mlpackage/Manifest.json", + sha256: "3513ae77df0cf9d4ccf658519ec7943c04b35edbd6500a996fea197666deadc7"), + Asset( + path: "part02.mlpackage/Data/com.apple.CoreML/model.mlmodel", + sha256: "0bcab81eeb64a8cb38f30ad5999f91da64291a7973431fabeb595978ddd30634"), + Asset( + path: "part02.mlpackage/Data/com.apple.CoreML/weights/weight.bin", + sha256: "0f1c235137345d7a8dd0bec7022bd906a90354d77ac80cd34ac2ea8308ec4994"), + Asset( + path: "part02.mlpackage/Manifest.json", + sha256: "10d60e47dd518d0c001e692330cbe4a7121f282e17dd870d3b3c82d035655d60"), + Asset( + path: "part03.mlpackage/Data/com.apple.CoreML/model.mlmodel", + sha256: "0bcab81eeb64a8cb38f30ad5999f91da64291a7973431fabeb595978ddd30634"), + Asset( + path: "part03.mlpackage/Data/com.apple.CoreML/weights/weight.bin", + sha256: "30c42923a00bba96303617bb3958a1cc3846cc46961eaba8153f3134bd9ea217"), + Asset( + path: "part03.mlpackage/Manifest.json", + sha256: "f6ac5ed4c39505a824c6e446177163f1484541516163b1b5000d611ae1d3b2f0"), + Asset( + path: "part04.mlpackage/Data/com.apple.CoreML/model.mlmodel", + sha256: "0bcab81eeb64a8cb38f30ad5999f91da64291a7973431fabeb595978ddd30634"), + Asset( + path: "part04.mlpackage/Data/com.apple.CoreML/weights/weight.bin", + sha256: "18588999cf26636cea57fbe27407dd686fa61eb132059d59a6aefbf0bbb668e3"), + Asset( + path: "part04.mlpackage/Manifest.json", + sha256: "2d0ac2e3d8906a2a865bada9a93debeb93791c6dbdd239494ab5b30c91866e61"), + Asset( + path: "part05.mlpackage/Data/com.apple.CoreML/model.mlmodel", + sha256: "0bcab81eeb64a8cb38f30ad5999f91da64291a7973431fabeb595978ddd30634"), + Asset( + path: "part05.mlpackage/Data/com.apple.CoreML/weights/weight.bin", + sha256: "d6bb092984a185b8c5a8034476199c7940b8fa0ef7ddd12407e94faa19340c31"), + Asset( + path: "part05.mlpackage/Manifest.json", + sha256: "6ebd70522296d5d8f5c486a10dfb598ba2e42ebb8e34da3ee3af8f298d738c71"), + Asset( + path: "part06.mlpackage/Data/com.apple.CoreML/model.mlmodel", + sha256: "0bcab81eeb64a8cb38f30ad5999f91da64291a7973431fabeb595978ddd30634"), + Asset( + path: "part06.mlpackage/Data/com.apple.CoreML/weights/weight.bin", + sha256: "12bdf4e71405cfbe7d4e536c3e254b69e9f0f72e92696266a39e21c13abbce73"), + Asset( + path: "part06.mlpackage/Manifest.json", + sha256: "acb70f8f49a6e777996a810ecf0fc6021389d973c91ec660f53b0f1d888f872a"), + Asset( + path: "part07.mlpackage/Data/com.apple.CoreML/model.mlmodel", + sha256: "f902c868ce2577fdbba3f23df4f1094de1a19afcea33c10ec54ed7c36af5853e"), + Asset( + path: "part07.mlpackage/Data/com.apple.CoreML/weights/weight.bin", + sha256: "2d24b9211e27e40d9ca4cf2feee4c76df38f4cd346cee99ef2f481068bc7d270"), + Asset( + path: "part07.mlpackage/Manifest.json", + sha256: "ccb89391958ae5b9f0969f60ee0bda953aacb41dd5ed0383d1a2b5cdc9d536ed"), + Asset(path: "pos_embed.f32", sha256: "8d123c332694df52f46467865e58923aa1fe47ee353780eb7a39467eea3490aa"), + Asset(path: "tokenizer.json", sha256: "5f9e4d4901a92b997e463c1f46055088b6cca5ca61a6522d1b9f64c4bb81cb42"), + Asset( + path: "VisionTower_P256.mlpackage/Data/com.apple.CoreML/model.mlmodel", + sha256: "4f5671928d0fac2ee2331b6b901ff519771097e392d235a122d92324c53226a6"), + Asset( + path: "VisionTower_P256.mlpackage/Data/com.apple.CoreML/weights/weight.bin", + sha256: "eaceac75cd1c45401c0f910577adb17a3a48a38b9557c335a0e51b7b9b5d0a9b"), + Asset( + path: "VisionTower_P256.mlpackage/Manifest.json", + sha256: "b0545cdc4546413a184efaff4723a21e4b23c8a973acaa9730109a403578b61d"), + Asset( + path: "VisionTower_P576.mlpackage/Data/com.apple.CoreML/model.mlmodel", + sha256: "50a8894c45ea7447a18d418245205657fcb83c65a9aae0658373393c7b8a94bf"), + Asset( + path: "VisionTower_P576.mlpackage/Data/com.apple.CoreML/weights/weight.bin", + sha256: "eaceac75cd1c45401c0f910577adb17a3a48a38b9557c335a0e51b7b9b5d0a9b"), + Asset( + path: "VisionTower_P576.mlpackage/Manifest.json", + sha256: "3c3053d6eef3e97ac79168f0316f1b97892c048bbe92ab47827031b30aa1b246"), + ] + + /// Ensure the bundle exists in the FluidUse cache and return its directory. + public static func ensure(cacheDirectory: URL? = nil, progress: Progress? = nil) async throws -> URL { + let root = cacheDirectory ?? LayaModelStore.defaultCacheDirectory() + let directory = root.appendingPathComponent("jev-9b-coreml", isDirectory: true) + let manager = FileManager.default + let verified = directory.appendingPathComponent(".verified-\(revision)") + if manager.fileExists(atPath: verified.path), + assets.allSatisfy({ manager.fileExists(atPath: directory.appendingPathComponent($0.path).path) }) + { + return directory + } + try manager.createDirectory(at: directory, withIntermediateDirectories: true) + for asset in assets { + try Task.checkCancellation() + let destination = directory.appendingPathComponent(asset.path) + if manager.fileExists(atPath: destination.path), + try ClefVisionModelStore.checksum(of: destination) == asset.sha256 + { + continue + } + try manager.createDirectory(at: destination.deletingLastPathComponent(), withIntermediateDirectories: true) + progress?(asset.path, 0) + let escaped = asset.path.addingPercentEncoding(withAllowedCharacters: .urlPathAllowed) ?? asset.path + guard let url = URL(string: "https://huggingface.co/\(repository)/resolve/\(revision)/\(escaped)") else { + throw JevError.invalidAsset("Invalid Hugging Face asset URL for \(asset.path)") + } + let (temporary, response) = try await URLSession.shared.download(from: url) + defer { try? manager.removeItem(at: temporary) } + guard let http = response as? HTTPURLResponse, http.statusCode == 200 else { + throw JevError.invalidAsset( + "HTTP \((response as? HTTPURLResponse)?.statusCode ?? -1) for \(asset.path)") + } + guard try ClefVisionModelStore.checksum(of: temporary) == asset.sha256 else { + throw JevError.invalidAsset("checksum mismatch: \(asset.path)") + } + if manager.fileExists(atPath: destination.path) { try manager.removeItem(at: destination) } + try manager.moveItem(at: temporary, to: destination) + progress?( + asset.path, Int64((try? manager.attributesOfItem(atPath: destination.path)[.size] as? Int64) ?? 0)) + } + manager.createFile(atPath: verified.path, contents: nil) + return directory + } +} diff --git a/Sources/FluidUse/Jev/JevTypes.swift b/Sources/FluidUse/Jev/JevTypes.swift new file mode 100644 index 0000000..fc14979 --- /dev/null +++ b/Sources/FluidUse/Jev/JevTypes.swift @@ -0,0 +1,219 @@ +import CoreGraphics +import Foundation + +/// A System 1 question kind, as JEV's `POST /v1/decide`. +public enum JevKind: Sendable, Equatable { + /// Yes / no: probabilities for `["false", "true"]`. + case noul + /// Ordered 0...5: probabilities for `["0", ..., "5"]`. + case score + /// Pick one of 2...256 options, in the given order. + case choice([String]) + + var name: String { + switch self { + case .noul: return "noul" + case .score: return "score" + case .choice: return "choice" + } + } +} + +/// One question about the shared state. +public struct JevQuestion: Sendable, Equatable { + public let kind: JevKind + public let text: String + + public init(_ kind: JevKind, _ text: String) { + self.kind = kind + self.text = text + } +} + +/// The state: text and images, in order (JEV's `state` list of strings and `{"image": ...}` parts). +public enum JevStatePart: Sendable { + case text(String) + case image(CGImage) +} + +/// Calibrated probabilities for one question. +public struct JevDecision: Sendable, Equatable { + public let options: [String] + public let probabilities: [Float] + + /// Most likely option. + public var choice: String { + guard let best = probabilities.indices.max(by: { probabilities[$0] < probabilities[$1] }) else { return "" } + return options[best] + } + /// P(true) for a `noul` question. + public var yes: Float? { options.firstIndex(of: "true").map { probabilities[$0] } } + /// Expected value for a `score` question. + public var expectedScore: Float { zip(probabilities.indices, probabilities).reduce(0) { $0 + Float($1.0) * $1.1 } } +} + +public struct JevResult: Sendable { + public let decisions: [JevDecision] + /// Prompt tokens of each question's row (the image is shared). + public let inputTokens: [Int] + public let visionMilliseconds: Double + public let decoderMilliseconds: Double + public let totalMilliseconds: Double +} + +public enum JevError: Error, LocalizedError, Sendable { + case invalidAsset(String) + case invalidInput(String) + case tooLong(tokens: Int, bucket: Int) + + public var errorDescription: String? { + switch self { + case .invalidAsset(let m): return "Invalid asset: \(m)" + case .invalidInput(let m): return "Invalid input: \(m)" + case .tooLong(let tokens, let bucket): return "Prompt needs \(tokens) tokens; the loaded bucket holds \(bucket)" + } + } +} + +/// The bundle's `config.json` and `head.json`. +struct JevConfig: Sendable { + let hidden: Int + let rotaryDim: Int + let ropeTheta: Double + let mropeSection: [Int] + let padID: Int + let vocab: Int + let buckets: [Int] + let imageTokenID: Int + let visionStartID: Int + let visionEndID: Int + let parts: [String] + let visionPackages: [(file: String, patches: Int)] + let posEmbed: String + let headRows: String + let headCount: Int + /// Image preprocessing and vision-tower host math are shared with clef-vision; `patches` is the largest bucket + /// (`clef(patches:)` gives the config for another one). + let clef: ClefVisionConfig + + // head.json + let headIDs: [Int] + let bias: [Float] + let ranges: [String: (Int, Int)] + let temperatures: [String: Float] + let labels: [String] + + init(manifest json: [String: Any], head: [String: Any]) throws { + func int(_ key: String, _ dict: [String: Any] = [:]) throws -> Int { + guard let v = ((dict.isEmpty ? json : dict)[key] as? NSNumber)?.intValue else { + throw JevError.invalidAsset("config.json: \(key)") + } + return v + } + guard let buckets = json["buckets"] as? [NSNumber], let parts = json["parts"] as? [String], + let theta = (json["rope_theta"] as? NSNumber)?.doubleValue, + let section = json["mrope_section"] as? [NSNumber], let vision = json["vision"] as? [String: Any], + let packages = vision["packages"] as? [[String: Any]], let posEmbed = vision["pos_embed"] as? String, + let image = json["image_processor"] as? [String: Any], let mean = image["image_mean"] as? [NSNumber], + let std = image["image_std"] as? [NSNumber], let headInfo = json["head"] as? [String: Any], + let rows = headInfo["rows"] as? String + else { throw JevError.invalidAsset("config.json is missing sections") } + hidden = try int("hidden_size") + rotaryDim = try int("rotary_dim") + ropeTheta = theta + mropeSection = section.map(\.intValue) + padID = try int("pad_id") + vocab = try int("vocab_size") + self.buckets = buckets.map(\.intValue).sorted() + imageTokenID = try int("image_token_id") + visionStartID = try int("vision_start_token_id") + visionEndID = try int("vision_end_token_id") + self.parts = parts + visionPackages = try packages.map { + guard let file = $0["file"] as? String else { throw JevError.invalidAsset("config.json: vision file") } + return (file, try int("patches", $0)) + }.sorted { $0.patches < $1.patches } + self.posEmbed = posEmbed + headRows = rows + headCount = try int("count", headInfo) + let headDim = try int("hidden_size", vision) / (try int("num_heads", vision)) + clef = ClefVisionConfig( + imageTokenID: imageTokenID, visionStartID: visionStartID, visionEndID: visionEndID, padID: padID, + vocabSize: vocab, hiddenSize: hidden, rotaryDim: rotaryDim, ropeTheta: theta, + mropeSection: mropeSection, lmBuckets: self.buckets, headLength: 0, headMaxQuestions: 0, + headMaxOptions: 0, + vision: .init( + patches: visionPackages.last?.patches ?? 0, patchSize: try int("patch_size", vision), + mergeSize: try int("spatial_merge_size", vision), temporalPatch: try int("temporal_patch_size", vision), + hidden: try int("hidden_size", vision), headDim: headDim, + depthPositions: try int("num_position_embeddings", vision), ropeTheta: 10_000), + image: .init( + minPixels: try int("min_pixels", image), maxPixels: try int("max_pixels", image), + mean: mean.map(\.floatValue), std: std.map(\.floatValue), + rescale: (image["rescale_factor"] as? NSNumber)?.floatValue ?? 1 / 255)) + + guard let ids = head["ids"] as? [NSNumber], let bias = head["bias"] as? [NSNumber], + let ranges = head["ranges"] as? [String: [NSNumber]], + let temps = head["temperatures"] as? [String: NSNumber], + let labels = head["labels"] as? [String] + else { throw JevError.invalidAsset("head.json is missing sections") } + headIDs = ids.map(\.intValue) + self.bias = bias.map(\.floatValue) + self.ranges = ranges.mapValues { ($0[0].intValue, $0[1].intValue) } + temperatures = temps.mapValues(\.floatValue) + self.labels = labels + guard headIDs.count == headCount else { throw JevError.invalidAsset("head.json ids vs head_rows count") } + } + + /// Rows of the head table and per-option bias for one question (`serve_decide.s1_pass`): noul and score read + /// their verbalizer slots; choice reads the trained A...P slots, then the remaining option labels with no bias. + func decisionRows(_ kind: JevKind) throws -> (rows: [Int], bias: [Float]) { + guard let (lo, hi) = ranges[kind.name] else { throw JevError.invalidAsset("head.json: no \(kind.name) range") } + guard case .choice(let options) = kind else { + return (Array(lo.. ClefVisionConfig { + let v = clef.vision + return ClefVisionConfig( + imageTokenID: clef.imageTokenID, visionStartID: clef.visionStartID, visionEndID: clef.visionEndID, + padID: clef.padID, vocabSize: clef.vocabSize, hiddenSize: clef.hiddenSize, rotaryDim: clef.rotaryDim, + ropeTheta: clef.ropeTheta, mropeSection: clef.mropeSection, lmBuckets: clef.lmBuckets, headLength: 0, + headMaxQuestions: 0, headMaxOptions: 0, + vision: .init( + patches: patches, patchSize: v.patchSize, mergeSize: v.mergeSize, temporalPatch: v.temporalPatch, + hidden: v.hidden, headDim: v.headDim, depthPositions: v.depthPositions, ropeTheta: v.ropeTheta), + image: clef.image) + } +} + +extension ClefVisionConfig { + init( + imageTokenID: Int, visionStartID: Int, visionEndID: Int, padID: Int, vocabSize: Int, hiddenSize: Int, + rotaryDim: Int, ropeTheta: Double, mropeSection: [Int], lmBuckets: [Int], headLength: Int, + headMaxQuestions: Int, headMaxOptions: Int, vision: VisionShape, image: ImageProcessing + ) { + self.imageTokenID = imageTokenID + self.visionStartID = visionStartID + self.visionEndID = visionEndID + self.padID = padID + self.vocabSize = vocabSize + self.hiddenSize = hiddenSize + self.rotaryDim = rotaryDim + self.ropeTheta = ropeTheta + self.mropeSection = mropeSection + self.lmBuckets = lmBuckets + self.headLength = headLength + self.headMaxQuestions = headMaxQuestions + self.headMaxOptions = headMaxOptions + self.vision = vision + self.image = image + } +} diff --git a/Sources/JevCheck/main.swift b/Sources/JevCheck/main.swift new file mode 100644 index 0000000..ad92720 --- /dev/null +++ b/Sources/JevCheck/main.swift @@ -0,0 +1,135 @@ +import CoreGraphics +import FluidUse +import Foundation +import ImageIO + +/// JEV-9B Core ML checks against the Python fixtures (model-lab `jev-9b-coreml/conversion/export_fixtures.py`). +/// +/// swift run -c release JevCheck parity # token ids, positions, probabilities +/// swift run -c release JevCheck latency # warm timings, one image + 3 questions +@main +@available(macOS 15.0, *) +struct JevCheck { + struct Fixtures: Decodable { + struct Record: Decodable { + let state: [Part] + let kind: String + let question: String + let options: [String]? + let inputIds: [Int] + let positionIds: [[Int]] + let probabilities: [Float] + } + enum Part: Decodable { + case text(String) + case image(String) + init(from decoder: Decoder) throws { + let c = try decoder.singleValueContainer() + if let s = try? c.decode(String.self) { + self = .text(s) + } else { + self = .image(try c.decode([String: String].self)["image"] ?? "") + } + } + } + let bucket: Int + let records: [Record] + } + + static func image(_ url: URL) throws -> CGImage { + guard let source = CGImageSourceCreateWithURL(url as CFURL, nil), + let image = CGImageSourceCreateImageAtIndex(source, 0, nil) + else { throw JevError.invalidInput("cannot read \(url.path)") } + return image + } + + static func question(_ r: Fixtures.Record) -> JevQuestion { + switch r.kind { + case "noul": return JevQuestion(.noul, r.question) + case "score": return JevQuestion(.score, r.question) + default: return JevQuestion(.choice(r.options ?? []), r.question) + } + } + + static func state(_ r: Fixtures.Record, in dir: URL) throws -> [JevStatePart] { + try r.state.map { + switch $0 { + case .text(let s): return .text(s) + case .image(let name): return .image(try image(dir.appendingPathComponent(name))) + } + } + } + + static func main() async throws { + let args = CommandLine.arguments + guard args.count >= 4 else { + print("usage: JevCheck parity|latency ") + exit(2) + } + let fixturesDir = URL(fileURLWithPath: args[2]) + let decoder = JSONDecoder() + decoder.keyDecodingStrategy = .convertFromSnakeCase + let fixtures = try decoder.decode( + Fixtures.self, from: Data(contentsOf: fixturesDir.appendingPathComponent("fixtures.json"))) + let loadStarted = Date() + let manager = try await JevManager.load(from: URL(fileURLWithPath: args[3]), bucket: fixtures.bucket) + print(String(format: "loaded in %.1f s", Date().timeIntervalSince(loadStarted))) + switch args[1] { + case "parity": + var idFailures = 0 + var flips = 0 + var maxDrift: Float = 0 + for (i, r) in fixtures.records.enumerated() { + let st = try state(r, in: fixturesDir) + let q = question(r) + let (ids, positions) = try await manager.tokens(state: st, question: q) + let idsOK = ids == r.inputIds && positions == r.positionIds + if !idsOK { idFailures += 1 } + let result = try await manager.decide(state: st, questions: [q]) + let p = result.decisions[0].probabilities + let drift = zip(p, r.probabilities).map { abs($0 - $1) }.max() ?? 0 + maxDrift = max(maxDrift, drift) + let argmax = { (v: [Float]) in v.indices.max { v[$0] < v[$1] }! } + let sorted = r.probabilities.sorted(by: >) + let nearTie = sorted.count > 1 && sorted[0] - sorted[1] < 0.02 // either answer is within noise + if argmax(p) != argmax(r.probabilities) && !nearTie { flips += 1 } + print( + String( + format: "%2d %-6@ tokens %3d %@ drift %.4f swift %@ python %@", i, r.kind as NSString, + ids.count, idsOK ? "ids OK " : "ids MISMATCH", drift, + p.map { String(format: "%.3f", $0) }.joined(separator: " "), + r.probabilities.map { String(format: "%.3f", $0) }.joined(separator: " "))) + } + print( + "records \(fixtures.records.count): token/position mismatches \(idFailures), argmax flips (python margin" + + " >= 0.02) \(flips), max prob drift \(maxDrift)") + if idFailures > 0 || flips > 0 { exit(1) } + case "latency": + guard + let r = fixtures.records.first(where: { + $0.state.contains { if case .image = $0 { true } else { false } } + }) + else { exit(1) } + let st = try state(r, in: fixturesDir) + let qs = [ + JevQuestion(.noul, "Is there an enemy in this picture?"), + JevQuestion(.noul, "Is there a pipe in this picture?"), + JevQuestion(.noul, "Is there a gap in this picture?"), + ] + _ = try await manager.decide(state: st, questions: qs) + var totals: [Double] = [] + for _ in 0..<10 { + let result = try await manager.decide(state: st, questions: qs) + totals.append(result.totalMilliseconds) + print( + String( + format: "vision %.0f ms decoder %.0f ms total %.0f ms", result.visionMilliseconds, + result.decoderMilliseconds, result.totalMilliseconds)) + } + print(String(format: "median %.0f ms for 1 image + 3 questions", totals.sorted()[totals.count / 2])) + default: + print("unknown mode \(args[1])") + exit(2) + } + } +} diff --git a/Sources/JevMarioDemo/ContentView.swift b/Sources/JevMarioDemo/ContentView.swift new file mode 100644 index 0000000..266a868 --- /dev/null +++ b/Sources/JevMarioDemo/ContentView.swift @@ -0,0 +1,150 @@ +import FluidUse +import SwiftUI + +@available(macOS 15.0, *) +struct ContentView: View { + @State private var player = MarioPlayer() + let modelDirectory: URL? // JEV-9B bundle; nil = download on first use + let bridgeCommand: [String] + + var body: some View { + HStack(alignment: .top, spacing: 24) { + VStack(alignment: .leading, spacing: 10) { + ZStack(alignment: .topTrailing) { + if let frame = player.frame { + Image(decorative: frame, scale: 1).resizable().interpolation(.none) + .aspectRatio(256 / 240, contentMode: .fit) + } else { + Rectangle().fill(.black).aspectRatio(256 / 240, contentMode: .fit) + } + if player.thinking { + Text("thinking…").font(.system(.body, design: .monospaced)).padding(6) + .background(.black.opacity(0.6)).foregroundStyle(.white).padding(8) + } + } + .frame(minWidth: 384, maxWidth: .infinity) + HStack { + Text(player.level.isEmpty ? "" : "World \(player.level)").font(.title3.bold()) + Spacer() + if !player.results.isEmpty { + Text("\(player.results.filter(\.flag).count) flags").font(.title3.monospacedDigit()) + } + } + LazyVGrid(columns: [GridItem(.adaptive(minimum: 40), spacing: 4)], alignment: .leading, spacing: 4) { + ForEach(Array(player.results.enumerated()), id: \.offset) { _, r in + Text(r.level).font(.system(.caption, design: .monospaced)) + .padding(.horizontal, 6).padding(.vertical, 2) + .background(r.flag ? Color.green.opacity(0.35) : Color.red.opacity(0.35), in: Capsule()) + } + } + ProgressView(value: Double(min(player.x, MarioPlayer.flagX)), total: Double(MarioPlayer.flagX)) + Text(player.status).foregroundStyle(.secondary) + } + VStack(alignment: .leading, spacing: 8) { + Picker("Model", selection: $player.brain) { + ForEach(MarioPlayer.Brain.allCases) { Text($0.rawValue).tag($0) } + } + .pickerStyle(.segmented).labelsHidden() + if player.brain == .tiny { + Text( + "Every move: a \(player.studentParameters.isEmpty ? "small" : player.studentParameters + "-parameter") model looks at the pixels around Mario" + ) + .font(.caption).foregroundStyle(.secondary) + if !player.studentPlacement.isEmpty { + Label(player.studentPlacement, systemImage: "cpu").font(.callout.bold()).foregroundStyle(.green) + } + Toggle("Lookahead: try moves on a save state first", isOn: $player.lookahead).toggleStyle(.switch) + .font(.callout).controlSize(.small) + if player.lookahead { + Text("\(player.simulated) moves simulated for this move · \(player.simulatedTotal) this level") + .font(.system(.caption, design: .monospaced)).foregroundStyle(.secondary) + } + LazyVGrid(columns: [GridItem(.flexible(), spacing: 14), GridItem(.flexible())], spacing: 6) { + ForEach(Array(player.studentMoves.enumerated()), id: \.offset) { i, move in + let p = i < player.moveProbabilities.count ? player.moveProbabilities[i] : 0 + HStack(spacing: 6) { + Text(move).font(.caption.bold()).frame(width: 56, alignment: .leading) + Bar(value: p, color: move == player.studentMove ? .orange : .gray) + Text(String(format: "%.2f", p)).font(.system(.caption2, design: .monospaced)) + } + } + } + } else { + Text("Every move: JEV-9B answers three yes/no questions about the screen just ahead of Mario") + .font(.callout).foregroundStyle(.secondary) + } + ForEach(player.brain == .jev ? MarioPlayer.questions : []) { q in + HStack(spacing: 12) { + Group { + if let image = player.viewImages[q.view] { + Image(decorative: image, scale: 1).resizable().interpolation(.none).scaledToFit() + } else { + Rectangle().fill(.quaternary) + } + } + .frame(width: 96, height: 96) + VStack(alignment: .leading, spacing: 4) { + Text(q.id).font(.headline) + let p = player.probabilities[q.id] ?? 0 + Bar(value: p, color: p > q.threshold ? .orange : .gray) + Text(String(format: "P(yes) %.2f", p)).font(.system(.caption, design: .monospaced)) + } + } + } + HStack(alignment: .firstTextBaseline) { + let shown = player.brain == .tiny ? player.studentMove : player.move.rawValue + Text(shown.uppercased()).font(.system(size: 24, weight: .bold, design: .rounded)) + .foregroundStyle(shown == "run" ? .green : .orange) + Spacer() + Text(String(format: player.lastMilliseconds < 10 ? "%.2f ms" : "%.0f ms", player.lastMilliseconds)) + .font(.system(.body, design: .monospaced)) + Text("· #\(player.decisions)").foregroundStyle(.secondary) + } + HStack { + Toggle( + MarioPlayer.speedCap.isFinite + ? String(format: "%.0f× speed", MarioPlayer.speedCap) : "Max speed", + isOn: $player.turbo + ).controlSize(.small).toggleStyle( + .switch) + Spacer() + Text(String(format: "%.0f fps · %.1f× real time", player.gameFPS, player.gameFPS / 60)) + .font(.system(.body, design: .monospaced)) + .foregroundStyle(player.turbo ? .orange : .secondary) + } + HStack { + Picker("Level", selection: $player.choice) { + ForEach(MarioPlayer.allLevels, id: \.self) { Text($0).tag($0) } + } + .frame(width: 130) + Button("Play") { player.restart(modelDirectory: modelDirectory, bridgeCommand: bridgeCommand) } + .buttonStyle(.borderedProminent) + Button("Play all \(MarioPlayer.playlist.count)") { + player.restart(modelDirectory: modelDirectory, bridgeCommand: bridgeCommand, all: true) + } + .buttonStyle(.borderedProminent).tint(.orange) + } + } + .frame(width: 340) + } + .padding(16) + .task { player.start(modelDirectory: modelDirectory, bridgeCommand: bridgeCommand) } + } +} + +/// A probability bar that redraws instantly (ProgressView animates and lags behind sub-second decisions). +struct Bar: View { + let value: Float + let color: Color + + var body: some View { + GeometryReader { geo in + ZStack(alignment: .leading) { + Capsule().fill(.quaternary) + Capsule().fill(color).frame(width: geo.size.width * CGFloat(min(max(value, 0), 1))) + } + } + .frame(height: 8) + .transaction { $0.animation = nil } + } +} diff --git a/Sources/JevMarioDemo/JevMarioDemoApp.swift b/Sources/JevMarioDemo/JevMarioDemoApp.swift new file mode 100644 index 0000000..fcb7b70 --- /dev/null +++ b/Sources/JevMarioDemo/JevMarioDemoApp.swift @@ -0,0 +1,41 @@ +import FluidUse +import SwiftUI + +/// JEV-9B or the small pixel student (Core ML) plays Super Mario Bros. +/// +/// swift run -c release JevMarioDemo # JEV-9B (downloaded on first use) +/// JEV_MARIO_BRAIN=student MARIO_STUDENT_MODEL=<.mlpackage> swift run -c release JevMarioDemo +/// +/// JEV-9B loads from `JEV_MODEL_DIR` or the pinned Hugging Face download (`JevModelStore`), only when picked. The +/// emulator runs in `nes_bridge.py` through `uv` (`JEV_MARIO_BRIDGE` overrides the command). +@main +@available(macOS 15.0, *) +struct JevMarioDemoApp: App { + init() { setvbuf(stdout, nil, _IOLBF, 0) } // the decision log is tailed live (demo.sh) + + var body: some Scene { + WindowGroup("Mario on Core ML") { + ContentView( + modelDirectory: ProcessInfo.processInfo.environment["JEV_MODEL_DIR"].flatMap { + $0.isEmpty ? nil : URL(fileURLWithPath: $0) + }, + bridgeCommand: Self.bridgeCommand) + } + .defaultSize(width: 900, height: 540) + .windowResizability(.contentMinSize) + } + + static var bridgeCommand: [String] { + if let custom = ProcessInfo.processInfo.environment["JEV_MARIO_BRIDGE"] { + return custom.split(separator: " ").map(String.init) + } + let script = URL(fileURLWithPath: #filePath).deletingLastPathComponent().appendingPathComponent( + "retro_bridge.py") + return [ + "uv", "run", "--no-project", "--python", "3.12", "--with", "stable-retro", "--with", + "gym-super-mario-bros==7.4.0", "--with", + "nes-py", + "--with", "gym==0.26.2", "--with", "numpy<2", "python", script.path, + ] + } +} diff --git a/Sources/JevMarioDemo/MarioBridge.swift b/Sources/JevMarioDemo/MarioBridge.swift new file mode 100644 index 0000000..80d58f9 --- /dev/null +++ b/Sources/JevMarioDemo/MarioBridge.swift @@ -0,0 +1,139 @@ +import CoreGraphics +import Foundation + +/// The emulator process (`nes_bridge.py`): one JSON request line in, one JSON line + a raw RGB frame out. +actor MarioBridge { + struct Step: Sendable { + let frame: CGImage + let pixels: Data // raw 256x240 RGB + let x: Int + let screenX: Int + let screenY: Int + let onGround: Bool + let done: Bool + let flag: Bool + let time: Int + let frames: Int // emulator frames this request advanced + let control: Bool // the player has control (false during pipe entry / area change) + let area: Int // RAM 0x750 & 0x7F + let transition: Bool // this request held input through a game-controlled transition (pipe) + } + + struct Meta: Decodable { + let x: Int + let screenX: Int + let screenY: Int + let onGround: Bool + let done: Bool + let flag: Bool + let time: Int + let frames: Int? + let images: Int? + let control: Bool? + let area: Int? + let transition: Bool? + } + + static let width = 256, height = 240 + private let process = Process() + private let input = Pipe() + private let output = Pipe() + + /// `command` is the full command line, e.g. `uv run --with gym-super-mario-bros==7.4.0 ... python nes_bridge.py`. + init(command: [String]) throws { + process.executableURL = URL(fileURLWithPath: "/usr/bin/env") + process.arguments = command + process.standardInput = input + process.standardOutput = output + process.standardError = FileHandle.standardError + try process.run() + } + + deinit { process.terminate() } + + func reset(level: String? = nil, idle: Int = 0) throws -> Step { + guard let level else { return try request(#"{"cmd":"reset","idle":\#(idle)}"#).0 } + return try request(#"{"cmd":"reset","level":"\#(level)","idle":\#(idle)}"#).0 + } + + func step(_ buttons: UInt8) throws -> Step { try request(#"{"cmd":"step","buttons":\#(buttons)}"#).0 } + + /// A whole move in one call (holding the last input while the game has control); returns the final frame and + /// the frame 4 before it. Max speed uses this: one request and one redraw per move. + func run(_ buttons: [UInt8]) throws -> (step: Step, before: Data) { + let list = buttons.map(String.init).joined(separator: ",") + let (step, before) = try request(#"{"cmd":"run","buttons":[\#(list)],"hold":true}"#) + return (step, before ?? step.pixels) + } + + /// Save state slot for the student's lookahead (one slot); the reply carries no image. + func save() throws { try command(#"{"cmd":"save"}"#) } + + func load() throws { try command(#"{"cmd":"load"}"#) } + + private var buffer = Data() + + /// Bytes from the bridge's stdout, read in chunks (one read per byte was the bottleneck at ~4 ms per call). + private func read(count: Int) throws -> Data { + while buffer.count < count { + let chunk = output.fileHandleForReading.availableData + guard !chunk.isEmpty else { throw BridgeError.closed } + buffer.append(chunk) + } + let out = buffer.prefix(count) + buffer.removeFirst(count) + return Data(out) + } + + private func readLine() throws -> Data { + while true { + if let newline = buffer.firstIndex(of: UInt8(ascii: "\n")) { + let line = buffer[buffer.startIndex.. Meta { + input.fileHandleForWriting.write(Data((line + "\n").utf8)) + let decoder = JSONDecoder() + decoder.keyDecodingStrategy = .convertFromSnakeCase + return try decoder.decode(Meta.self, from: readLine()) + } + + private func command(_ line: String) throws { _ = try meta(line) } + + private func request(_ line: String) throws -> (Step, Data?) { + let meta = try meta(line) + let size = Self.width * Self.height * 3 + func readImage() throws -> Data { try read(count: size) } + let pixels = try readImage() + let before = (meta.images ?? 1) > 1 ? try readImage() : nil + guard let provider = CGDataProvider(data: pixels as CFData), + let frame = CGImage( + width: Self.width, height: Self.height, bitsPerComponent: 8, bitsPerPixel: 24, + bytesPerRow: Self.width * 3, space: CGColorSpaceCreateDeviceRGB(), + bitmapInfo: CGBitmapInfo(rawValue: CGImageAlphaInfo.none.rawValue), provider: provider, decode: nil, + shouldInterpolate: false, intent: .defaultIntent) + else { throw BridgeError.badFrame } + let step = Step( + frame: frame, pixels: pixels, x: meta.x, screenX: meta.screenX, screenY: meta.screenY, + onGround: meta.onGround, done: meta.done, flag: meta.flag, time: meta.time, frames: meta.frames ?? 1, + control: meta.control ?? true, area: meta.area ?? 0, transition: meta.transition ?? false) + return (step, before) + } + + enum BridgeError: Error, LocalizedError { + case closed, badFrame + var errorDescription: String? { + switch self { + case .closed: return "The emulator bridge exited (is gym-super-mario-bros installed? see README)" + case .badFrame: return "Could not decode an emulator frame" + } + } + } +} diff --git a/Sources/JevMarioDemo/MarioPlayer.swift b/Sources/JevMarioDemo/MarioPlayer.swift new file mode 100644 index 0000000..f00e4b4 --- /dev/null +++ b/Sources/JevMarioDemo/MarioPlayer.swift @@ -0,0 +1,447 @@ +import CoreGraphics +import FluidUse +import Foundation +import Observation + +/// The control loop. While Mario is on the ground the emulator pauses for a decision, then the move plays out frame by +/// frame at 60 fps. +/// - JEV-9B: three yes/no questions about close-ups just ahead of him (enemy? pipe or step? gap?); any yes -> long +/// running jump, else run. Mirrors model-lab `jev-9b-coreml/conversion/mario.py`. +/// - Student: a small pixel model (~9M parameters) picks run / jump / hop / wait from the pixels around Mario +/// (model-lab `mario-tiny`). +@available(macOS 15.0, *) +@MainActor +@Observable +final class MarioPlayer { + struct Question: Identifiable { + let id: String // panel label + let view: String + let text: String + let threshold: Float + } + + enum Move: String { case run, jump, hop, wait } + + enum Brain: String, CaseIterable, Identifiable { + case jev = "JEV-9B · 3 questions" + case tiny = "Student · pixels" + var id: String { rawValue } + } + + /// (x offset from Mario, width, top relative to his feet, height) in screen px, and the state text. + static let views: [String: (dx: Int, width: Int, dy: Int, height: Int, state: String)] = [ + "ahead": (16, 64, -48, 80, "Close-up of the area just ahead of Mario in a Super Mario Bros. level:"), + "above": ( + 16, 64, -48, 48, + "Close-up of the space just ahead of Mario in a Super Mario Bros. level, at the height of his body:" + ), + "ground": (16, 96, -32, 64, "Close-up of the ground in a Super Mario Bros. level:"), + ] + static let questions: [Question] = [ + .init( + id: "enemy ahead", view: "ahead", text: "Is there an enemy (a brown Goomba or a Koopa) in this picture?", + threshold: 0.5), + .init( + id: "pipe / step ahead", view: "above", + text: "Is there a solid obstacle (a pipe, a block or a step) in this picture?", threshold: 0.5), + .init( + id: "gap ahead", view: "ground", text: "Is the brick floor in this picture interrupted by a gap?", + threshold: 0.5), + ] + + static let right: UInt8 = 0b1000_0000, a: UInt8 = 0b0000_0001, b: UInt8 = 0b0000_0010 + static let moves: [Move: [UInt8]] = [ + .run: Array(repeating: right | b, count: 6), + .jump: [right | b] + Array(repeating: right | a | b, count: 22) + Array(repeating: right | b, count: 4), + .hop: [right | b] + Array(repeating: right | a | b, count: 8) + Array(repeating: right | b, count: 4), + .wait: Array(repeating: 0, count: 6), + ] + static let flagX = 3160 + /// Optional cap on max speed, in multiples of real time (MARIO_SPEED_CAP); none by default. + static let speedCap = Double(ProcessInfo.processInfo.environment["MARIO_SPEED_CAP"] ?? "") ?? .infinity + + // published state + var frame: CGImage? + var viewImages: [String: CGImage] = [:] + var probabilities: [String: Float] = [:] + var move: Move = .run + var thinking = false + var decisions = 0 + var lastMilliseconds: Double = 0 + var x = 40 + var status = "Loading…" + var finished = false + var brain: Brain = ProcessInfo.processInfo.environment["JEV_MARIO_BRAIN"] == "student" ? .tiny : .jev + var moveProbabilities: [Float] = [] // student: per studentMoves + var studentMoves: [String] = [] + var studentMove = "" + var studentParameters = "" + /// Max speed: no 60 fps pacing, the emulator runs as fast as the bridge and the model allow. + /// MARIO_TURBO=1 starts with the switch on (default off). + var turbo = ProcessInfo.processInfo.environment["MARIO_TURBO"] == "1" + var gameFPS: Double = 0 // emulator frames per second over the last second + private var frameLog: [(Date, Int)] = [] + private var levelFrames = 0 + var studentPlacement = "" + /// Student lookahead (model-lab mario-tiny `student_retro.py play --lookahead 6`): candidate moves are simulated on + /// an emulator save state before each real move. MARIO_LOOKAHEAD=0 starts with it off. + var lookahead = ProcessInfo.processInfo.environment["MARIO_LOOKAHEAD"] != "0" + var simulated = 0 // moves simulated for the last real move + var simulatedTotal = 0 + static let lookaheadDepth = 6 + + private var manager: JevManager? + private var tiny: MarioStudentPolicy? + private var history: [Data] = [] // last 5 frames' pixels (the student also sees 4 frames back) + private var bridge: MarioBridge? + private var loop: Task? + private var started = false + + /// Starts the first game; SwiftUI may re-run the view's `.task`, so later calls are ignored (use `restart`). + /// Loads the model and the emulator at launch; nothing plays until Play. + func start(modelDirectory: URL?, bridgeCommand: [String]) { + guard !started else { return } + started = true + run(modelDirectory: modelDirectory, bridgeCommand: bridgeCommand, playing: false) + } + + private func run(modelDirectory: URL?, bridgeCommand: [String], playing: Bool = true) { + let previous = loop + previous?.cancel() + loop = Task { + await previous?.value // the old game stops before this one resets the emulator + do { + let loadStarted = Date() + switch brain { + case .jev where manager == nil: + status = + modelDirectory == nil + ? "Downloading JEV-9B Core ML (~13 GB, first use only)…" : "Loading JEV-9B…" + let directory: URL + if let modelDirectory { + directory = modelDirectory + } else { + directory = try await JevModelStore.ensure() + } + status = "Loading JEV-9B (Core ML)…" + manager = try await JevManager.load(from: directory, bucket: 256) + case .tiny where tiny == nil: + guard let path = ProcessInfo.processInfo.environment["MARIO_STUDENT_MODEL"] else { + throw DemoError.noStudentModel + } + let policy = try await MarioStudentPolicy(package: URL(fileURLWithPath: path)) + studentParameters = policy.parameterLabel + studentMoves = policy.moves + studentPlacement = policy.placement + print("student: \(policy.parameterLabel) params, \(policy.placement)") + tiny = policy + default: break + } + status = String(format: "\(brain.rawValue) ready in %.1f s", Date().timeIntervalSince(loadStarted)) + if bridge == nil { bridge = try MarioBridge(command: bridgeCommand) } + // MARIO_AUTOPLAY=1 (benchmarks only, off by default): press Play once after loading + let autoplay = !playing && ProcessInfo.processInfo.environment["MARIO_AUTOPLAY"] == "1" + guard playing || autoplay else { + finished = true + status += " · pick a level and press Play" + return + } + try await play() + } catch is CancellationError { + } catch { + status = "Error: \(error.localizedDescription)" + } + } + } + + enum DemoError: Error, LocalizedError { + case noStudentModel + var errorDescription: String? { "Set MARIO_STUDENT_MODEL to the student's .mlpackage path" } + } + + func restart(modelDirectory: URL?, bridgeCommand: [String], all: Bool = false) { + playAll = all + finished = false + decisions = 0 + run(modelDirectory: modelDirectory, bridgeCommand: bridgeCommand) + } + + var level = "" + var results: [(level: String, flag: Bool)] = [] + /// The level Play starts (one level per press). + var choice = ProcessInfo.processInfo.environment["MARIO_LEVEL"].flatMap { $0.isEmpty ? nil : $0 } ?? "1-1" + static let allLevels = (1...8).flatMap { w in (1...4).map { "\(w)-\($0)" } } + /// "Play all": the levels the student + lookahead clears from the start every time (deterministic emulator; model-lab + /// mario-tiny r10, idle 0). MARIO_PLAYLIST="1-1,6-1,..." overrides. + static let playlist = + ProcessInfo.processInfo.environment["MARIO_PLAYLIST"].flatMap { + $0.isEmpty ? nil : $0.split(separator: ",").map(String.init) + } + ?? [ + "1-1", "1-2", "1-3", "1-4", "2-3", "2-4", "3-1", "3-2", "3-3", "3-4", "4-1", "4-2", "5-2", "5-3", "5-4", + "6-1", "6-2", "6-3", "6-4", "7-1", "7-3", "8-1", "8-2", + ] + private var playAll = false + var runStarted: Date? + var runSeconds: Double = 0 + + private func play() async throws { + let levels = playAll ? Self.playlist : [choice] + if playAll { results = [] } + let started = Date() + runStarted = started + for (index, next) in levels.enumerated() { + level = next + decisions = 0 + simulatedTotal = 0 + let flag = try await playLevel(next, idle: 0) + results.append((next, flag)) + runSeconds = Date().timeIntervalSince(started) + print("LEVEL \(next): \(status)") + if index < levels.count - 1 { try await Task.sleep(for: .seconds(1)) } // a beat on the flag + } + if playAll { + let flags = results.filter(\.flag).count + status = String(format: "%d/%d levels reached the flag in %.0f s", flags, levels.count, runSeconds) + print("RUN: \(status)") + } + finished = true + } + + private func playLevel(_ level: String, idle: Int) async throws -> Bool { + guard let bridge else { return false } + var step = try await bridge.reset(level: level, idle: idle) + levelFrames = 0 + let levelStarted = Date() + defer { + let seconds = Date().timeIntervalSince(levelStarted) + print( + String( + format: "level %@: %d frames in %.1f s = %.0f fps (%.1fx real time)", level, levelFrames, + seconds, Double(levelFrames) / seconds, Double(levelFrames) / seconds / 60)) + } + show(step) + history = Array(repeating: step.pixels, count: 5) + while !step.done { + try Task.checkCancellation() + // 1. decide (emulator paused) + if brain == .tiny, let tiny { + let started = Date() + moveProbabilities = try tiny.decide(now: step.pixels, before: history[0], marioX: step.screenX) + lastMilliseconds = Date().timeIntervalSince(started) * 1000 + let ranked = moveProbabilities.indices.sorted { moveProbabilities[$0] > moveProbabilities[$1] } + .map { tiny.moves[$0] } + studentMove = ranked.first ?? "run" + decisions += 1 + print(String(format: "%3d x %4d %@ %.2f ms", decisions, step.x, studentMove, lastMilliseconds)) + if tiny.everyMacro && lookahead { + studentMove = try await look(from: step, ranked: ranked, tiny: tiny, bridge: bridge) + simulatedTotal += simulated + } + if tiny.everyMacro { + step = try await performMacro( + MarioStudentPolicy.solverMacros[studentMove] ?? [], from: step, bridge: bridge) + } else { + move = Move(rawValue: studentMove) ?? .run + step = try await perform(move, from: step, bridge: bridge) + } + continue + } + guard let manager else { return false } + thinking = true + let started = Date() + var probs: [String: Float] = [:] + var images: [String: CGImage] = [:] + for (name, spec) in Self.views { + guard let image = Self.view(step.frame, x: step.screenX, y: step.screenY, spec: spec) else { continue } + images[name] = image + let asked = Self.questions.filter { $0.view == name } + let result = try await manager.decide( + state: [.text(spec.state), .image(image)], questions: asked.map { JevQuestion(.noul, $0.text) }) + for (q, d) in zip(asked, result.decisions) { probs[q.id] = d.yes ?? 0 } + } + lastMilliseconds = Date().timeIntervalSince(started) * 1000 + viewImages = images + probabilities = probs + move = Self.questions.contains { (probs[$0.id] ?? 0) > $0.threshold } ? .jump : .run + decisions += 1 + thinking = false + let line = Self.questions.map { String(format: "%@ %.2f", $0.id, probs[$0.id] ?? 0) }.joined( + separator: " ") + print( + String(format: "%3d x %4d %@ %@ %.0f ms", decisions, step.x, move.rawValue, line, lastMilliseconds)) + step = try await perform(move, from: step, bridge: bridge) + } + status = + step.flag + ? "Reached the flag in \(decisions) decisions" : "Mario died at x \(step.x) after \(decisions) decisions" + if brain == .tiny && lookahead { status += " · \(simulatedTotal) moves simulated" } + return step.flag + } + + /// The first candidate (in the student's order) that, simulated on a save state with the student's next 6 moves, + /// survives and gains ground: x +16 px, another area, or a pipe (entering a side pipe drops x). If the student's + /// continuation stalls, the same move repeated for 18 moves is tried too (a swim stroke only rises if repeated). + /// Else the candidate that lasts longest, then gets furthest. Restores the save state before returning. + private func look( + from step: MarioBridge.Step, ranked: [String], tiny: MarioStudentPolicy, bridge: MarioBridge + ) async throws -> String { + try await bridge.save() + simulated = 0 + var best = ranked.first ?? "run" + var bestScore = (-1, -1) + for move in ranked { + for repeating in [false, true] { + try Task.checkCancellation() + try await bridge.load() + let horizon = repeating ? 3 * Self.lookaheadDepth : Self.lookaheadDepth + var (sim, before) = try await bridge.run(MarioStudentPolicy.solverMacros[move] ?? []) + var transition = sim.transition + var lasted = 0 + simulated += 1 + while !sim.done && lasted < horizon { + var next = move + if !repeating { + let started = Date() + let p = try tiny.decide(now: sim.pixels, before: before, marioX: sim.screenX) + print(String(format: " sim %.2f ms", Date().timeIntervalSince(started) * 1000)) + next = tiny.moves[p.indices.max { p[$0] < p[$1] } ?? 0] + } + (sim, before) = try await bridge.run(MarioStudentPolicy.solverMacros[next] ?? []) + transition = transition || sim.transition + lasted += 1 + simulated += 1 + } + let dead = sim.done && !sim.flag + let moved = sim.x > step.x + 16 || sim.area != step.area || transition + if sim.flag || (!sim.done && lasted >= horizon && moved) { + try await bridge.load() + return move + } + let score = (dead && lasted == 0 ? -1 : lasted, sim.x < 0xF000 ? sim.x : 0) + if score > bestScore { + best = move + bestScore = score + } + } + } + try await bridge.load() + return best + } + + /// A solver-style 8-frame macro. Max speed: one batched bridge call and one redraw per macro. Otherwise frame by + /// frame at 60 fps, holding the last input while the game has control (pipe entry, area change). + private func performMacro( + _ buttons: [UInt8], from start: MarioBridge.Step, bridge: MarioBridge + ) async throws + -> MarioBridge.Step + { + guard !buttons.isEmpty else { return start } + try Task.checkCancellation() + if turbo { + let tick = Date() + let (step, before) = try await bridge.run(buttons) + show(step) + history = [before, before, before, before, step.pixels] + countFrames(step.frames) + // cap game time at speedCap x real time (60 fps x cap); uncapped by default + let budget = Double(step.frames) / (60 * Self.speedCap) + let elapsed = Date().timeIntervalSince(tick) + if elapsed < budget { try await Task.sleep(for: .seconds(budget - elapsed)) } + return step + } + var step = start + var index = 0 + var held = 0 + while !step.done { + let press: UInt8 + if index < buttons.count { + press = buttons[index] + index += 1 + } else if !step.control && held < 300 { + press = buttons[buttons.count - 1] + held += 1 + } else { + break + } + let tick = Date() + try Task.checkCancellation() + step = try await bridge.step(press) + show(step) + history = Array(history.dropFirst()) + [step.pixels] + countFrames(1) + let elapsed = Date().timeIntervalSince(tick) + if elapsed < 1 / 60 { try await Task.sleep(for: .seconds(1 / 60 - elapsed)) } + } + return step + } + + /// Emulator frames per second over the last second (counts batched frames too). + private func countFrames(_ n: Int) { + levelFrames += n + let now = Date() + frameLog.append((now, n)) + frameLog.removeAll { now.timeIntervalSince($0.0) > 1 } + gameFPS = Double(frameLog.reduce(0) { $0 + $1.1 }) + } + + /// Play the move at 60 fps, then keep running until Mario lands (max 90 frames). + private func perform( + _ move: Move, from start: MarioBridge.Step, bridge: MarioBridge + ) async throws + -> MarioBridge.Step + { + var step = start + let buttons = Self.moves[move] ?? [] + var index = 0 + var airborne = 0 + while !step.done { + let press: UInt8 + if index < buttons.count { + press = buttons[index] + index += 1 + } else if !step.onGround && airborne < 90 { + press = Self.right | Self.b + airborne += 1 + } else { + break + } + let tick = Date() + try Task.checkCancellation() + step = try await bridge.step(press) + show(step) + history = Array(history.dropFirst()) + [step.pixels] + countFrames(1) + let elapsed = Date().timeIntervalSince(tick) + if !turbo, elapsed < 1 / 60 { try await Task.sleep(for: .seconds(1 / 60 - elapsed)) } + } + return step + } + + private func show(_ step: MarioBridge.Step) { + frame = step.frame + x = step.x + } + + /// A close-up: Mario's feet are 32 px below the RAM y; upscaled 4x with nearest-neighbour (PIL NEAREST). + static func view( + _ frame: CGImage, x: Int, y: Int, spec: (dx: Int, width: Int, dy: Int, height: Int, state: String) + ) + -> CGImage? + { + let x0 = min(max(x + spec.dx, 0), frame.width - spec.width) + let y0 = min(max(y + 32 + spec.dy, 0), frame.height - spec.height) + guard let crop = frame.cropping(to: CGRect(x: x0, y: y0, width: spec.width, height: spec.height)) else { + return nil + } + let scale = 4 + guard + let context = CGContext( + data: nil, width: spec.width * scale, height: spec.height * scale, bitsPerComponent: 8, bytesPerRow: 0, + space: CGColorSpaceCreateDeviceRGB(), bitmapInfo: CGImageAlphaInfo.noneSkipLast.rawValue) + else { return nil } + context.interpolationQuality = .none + context.draw(crop, in: CGRect(x: 0, y: 0, width: spec.width * scale, height: spec.height * scale)) + return context.makeImage() + } +} diff --git a/Sources/JevMarioDemo/MarioStudentPolicy.swift b/Sources/JevMarioDemo/MarioStudentPolicy.swift new file mode 100644 index 0000000..ffcd098 --- /dev/null +++ b/Sources/JevMarioDemo/MarioStudentPolicy.swift @@ -0,0 +1,106 @@ +@preconcurrency import CoreML +import Foundation + +/// The pixel student (model-lab `mario-tiny`, a ~9M-parameter MobileNetV3 trained with RAM tile / enemy targets): +/// the screen around Mario now and 4 frames earlier -> logits over its moves (named in the package metadata; the first +/// `moves.count` outputs). A 4-move model (run / jump / hop / wait) decides when Mario lands; an 11-move model +/// (trained on the level solver) decides every 8 frames. Under 1 ms per decision on the Neural Engine. +@available(macOS 15.0, *) +@MainActor +final class MarioStudentPolicy { + /// Move names in output order, from the package metadata. + let moves: [String] + /// Solver-style model: 8-frame macros, one decision per macro. + var everyMacro: Bool { moves.count > 4 } + + static let right: UInt8 = 0x80, left: UInt8 = 0x40, down: UInt8 = 0x20, a: UInt8 = 0x01, b: UInt8 = 0x02 + /// The level solver's macros (model-lab mario-tiny/solver.py), 8 frames each. + static let solverMacros: [String: [UInt8]] = [ + "run": Array(repeating: right | b, count: 8), + "walk": Array(repeating: right, count: 8), + "hold": Array(repeating: right | b | a, count: 8), + "jump": [right | b] + Array(repeating: right | b | a, count: 7), + "hop": [right | b] + Array(repeating: right | b | a, count: 2) + Array(repeating: right | b, count: 5), + "up": [0] + Array(repeating: a, count: 7), + "left": Array(repeating: left | b, count: 8), + "leftjump": [left] + Array(repeating: left | a, count: 7), + "wait": Array(repeating: 0, count: 8), + "swim": Array(repeating: right | a, count: 2) + Array(repeating: right, count: 6), + "down": Array(repeating: down, count: 8), + ] + static let cropLeft = 32, cropWidth = 192, scale = 2 // 2 blocks behind Mario to 10 ahead, half resolution + private let model: MLModel + /// Parameter count from the package metadata (export_coreml.py), e.g. "9.0M". + let parameterLabel: String + /// Where Core ML runs the model's ops (MLComputePlan), e.g. "Neural Engine · 185/185 ops". + private(set) var placement = "" + private let compiledURL: URL + + init(package: URL) async throws { + let compiled = try await MLModel.compileModel(at: package) + compiledURL = compiled + let configuration = MLModelConfiguration() + configuration.computeUnits = .cpuAndNeuralEngine + model = try await MLModel.load(contentsOf: compiled, configuration: configuration) + let metadata = model.modelDescription.metadata[.creatorDefinedKey] as? [String: String] + let count = Double(metadata?["parameters"] ?? "") ?? 0 + parameterLabel = count >= 1e6 ? String(format: "%.1fM", count / 1e6) : String(format: "%.0fK", count / 1e3) + moves = (metadata?["moves"] ?? "run,jump,hop,wait").split(separator: ",").map(String.init) + placement = await Self.placement(of: compiled) + } + + /// Ops per preferred compute device from Core ML's compute plan (activity meters do not see the Neural Engine). + static func placement(of compiled: URL) async -> String { + let configuration = MLModelConfiguration() + configuration.computeUnits = .cpuAndNeuralEngine + guard let plan = try? await MLComputePlan.load(contentsOf: compiled, configuration: configuration), + case .program(let program) = plan.modelStructure, let main = program.functions["main"] + else { return "" } + var ane = 0 + var gpu = 0 + var cpu = 0 + for op in main.block.operations where op.operatorName != "const" { + switch plan.deviceUsage(for: op)?.preferred { + case .neuralEngine: ane += 1 + case .gpu: gpu += 1 + case .cpu: cpu += 1 + default: break + } + } + let total = ane + gpu + cpu + guard total > 0 else { return "" } + if ane == total { return "Neural Engine · \(ane)/\(total) ops" } + return "Neural Engine \(ane) · GPU \(gpu) · CPU \(cpu) ops" + } + + /// Probabilities per move in `moves` order. `now` / `before` are raw 256x240 RGB frames. + func decide(now: Data, before: Data, marioX: Int) throws -> [Float] { + let height = 240 / Self.scale + let width = Self.cropWidth / Self.scale + let input = try MLMultiArray(shape: [1, 6, height as NSNumber, width as NSNumber], dataType: .float16) + let out = input.dataPointer.assumingMemoryBound(to: Float16.self) + for (slot, frame) in [now, before].enumerated() { + frame.withUnsafeBytes { raw in + let px = raw.bindMemory(to: UInt8.self) + for y in 0..= 0 && sx < 256 + for c in 0..<3 { + let value: Float16 = inside ? Float16(Float(px[(sy * 256 + sx) * 3 + c]) / 255) : 0 + out[((slot * 3 + c) * height + y) * width + x] = value + } + } + } + } + } + let result = try model.prediction(from: MLDictionaryFeatureProvider(dictionary: ["crop": input])) + guard let logits = result.featureValue(for: "logits")?.multiArrayValue else { return [] } + let values = (0.. Sources/JevMarioDemo/demo.sh the small pixel student +# MONITOR=ane ... bars also show Neural Engine power via `sudo powermetrics` (type your password in the top pane; +# asitop 0.0.24 crashes on M5 Pro: no E0-Cluster in powermetrics) +# The emulator bridge runs through uv with gym-super-mario-bros (it bundles the NES ROM; local use only). +set -e +ROOT=$(cd "$(dirname "$0")/../.." && pwd) +BUNDLE=${1:-${JEV_MODEL_DIR:-}} +LOG=${TMPDIR:-/tmp}/jev-mario-demo.log + +[[ -z "$BUNDLE" || -f "$BUNDLE/config.json" ]] || { echo "no jev-9b-coreml bundle at $BUNDLE"; exit 1; } +[[ "${JEV_MARIO_BRAIN:-}" != student || -d "${MARIO_STUDENT_MODEL:-}" ]] || { echo "set MARIO_STUDENT_MODEL to the student .mlpackage"; exit 1; } +command -v uv >/dev/null || { echo "uv is required for the emulator bridge (https://docs.astral.sh/uv/)"; exit 1; } +swift build -c release --product JevMarioDemo --package-path "$ROOT" +: > "$LOG" +# a session left from an earlier run may have lost its Ghostty window: start fresh so the window always opens +tmux kill-session -t jev-mario 2>/dev/null || true +if true; then + TOP="python3 '$ROOT/Sources/ClefFlashDemo/bars.py' --log '$LOG'" + [[ "${MONITOR:-}" == ane ]] && TOP="$TOP --ane" || true + tmux new-session -d -s jev-mario -x 160 -y 60 "$TOP" + tmux set-option -t jev-mario remain-on-exit on + tmux set-option -t jev-mario window-size latest # follow the Ghostty window, not the detached 160-column default + tmux split-window -h -t jev-mario "tail -n 300 -F '$LOG'" # bars left, decision log right + BEFORE=$(pgrep -n -x ghostty || true) + open -na /Applications/Ghostty.app --args -e "$(command -v tmux)" attach -t jev-mario + for _ in {1..20}; do # the new Ghostty process can take a few seconds to appear + GHOSTTY_PID=$(pgrep -n -x ghostty || true) + [[ -n "$GHOSTTY_PID" && "$GHOSTTY_PID" != "$BEFORE" ]] && break + sleep 0.5 + done +fi +pkill -x JevMarioDemo 2>/dev/null || true +MARIO_PLAYLIST="${MARIO_PLAYLIST:-}" MARIO_TURBO="${MARIO_TURBO:-}" MARIO_LOOKAHEAD="${MARIO_LOOKAHEAD:-}" MARIO_LEVEL="${MARIO_LEVEL:-1-1}" JEV_MODEL_DIR="$BUNDLE" JEV_MARIO_BRAIN="${JEV_MARIO_BRAIN:-}" MARIO_STUDENT_MODEL="${MARIO_STUDENT_MODEL:-}" \ + "$ROOT/.build/release/JevMarioDemo" > "$LOG" 2> "$LOG.stderr" & +APP_PID=$! +# T layout: the game window centred across the top, the terminal full width underneath (bars | log) +( + for _ in {1..40}; do + osascript -e "tell application \"System Events\" to count windows of (first process whose unix id is $APP_PID)" 2>/dev/null \ + | grep -q '^[1-9]' && break + sleep 0.5 + done + read -r _ _ SCREEN_W SCREEN_H < <(osascript -e 'tell application "Finder" to get bounds of window of desktop' | tr -d ',') + MENU=33 + GAME_BOTTOM=$(osascript -e "tell application \"System Events\" to tell front window of (first process whose unix id is $APP_PID)" \ + -e "set {w, h} to size" -e "set position to {($SCREEN_W - w) div 2, $MENU}" -e "return $MENU + h" -e "end tell" 2>/dev/null) + [[ -n "${GHOSTTY_PID:-}" && -n "$GAME_BOTTOM" ]] || exit 0 + STUCK=0 + for _ in {1..40}; do # Ghostty re-places its window after opening; keep re-applying until it held for ~3 s + WIDTH=$(osascript -e "tell application \"System Events\" to tell front window of (first process whose unix id is $GHOSTTY_PID)" \ + -e "set {position, size} to {{0, $GAME_BOTTOM}, {$SCREEN_W, $SCREEN_H - $GAME_BOTTOM}}" -e "return item 1 of (get size)" \ + -e "end tell" 2>/dev/null) + [[ "$WIDTH" == "$SCREEN_W" ]] && STUCK=$((STUCK + 1)) || STUCK=0 + (( STUCK >= 6 )) && break + sleep 0.5 + done + osascript -e "tell application \"System Events\" to set frontmost of (first process whose unix id is $APP_PID) to true" 2>/dev/null +) & +echo "demo pid $APP_PID · log $LOG · first move after ~40 s of model loading (plus a ~13 GB download on first run)" diff --git a/Sources/JevMarioDemo/nes_bridge.py b/Sources/JevMarioDemo/nes_bridge.py new file mode 100644 index 0000000..fe6ae08 --- /dev/null +++ b/Sources/JevMarioDemo/nes_bridge.py @@ -0,0 +1,56 @@ +"""Super Mario Bros. emulator bridge for JevMarioDemo: frames in, button presses out, over stdin/stdout. + +Requests (one JSON line each): {"cmd": "reset", "level": "1-3", "idle": 17} (both optional: the current level, 0 idle +frames before play) or {"cmd": "step", "buttons": }. +Reply: one JSON line {"x", "screen_x", "screen_y", "on_ground", "done", "flag", "time"} followed by the frame as +240*256*3 raw RGB bytes. The decisions are made by the Swift app (JEV-9B on Core ML); this only emulates. + +MARIO_LEVEL picks the level (default 1-1). Needs gym-super-mario-bros 7.4.0 (it bundles the NES ROM; install it yourself, local use only), nes-py, gym 0.26.2, +numpy<2 -- `demo.sh` runs it through `uv run --with ...`. +""" +import json +import os +import sys +import warnings + +warnings.filterwarnings("ignore") +from gym_super_mario_bros import SuperMarioBrosEnv # noqa: E402 + +FLOAT_STATE, SCREEN_Y = 0x001D, 0x03B8 + + +def reply(env, obs, done, info): + meta = {"x": int(info.get("x_pos", 40)), "screen_x": int(env._left_x_position), "screen_y": int(env.ram[SCREEN_Y]), + "on_ground": bool(env.ram[FLOAT_STATE] == 0), "done": bool(done), "flag": bool(info.get("flag_get", False)), + "time": int(info.get("time", 400))} + sys.stdout.buffer.write((json.dumps(meta) + "\n").encode()) + sys.stdout.buffer.write(obs.tobytes()) + sys.stdout.buffer.flush() + + +def make(level: str): + world, stage = (int(v) for v in level.split("-")) + return SuperMarioBrosEnv(rom_mode="vanilla", lost_levels=False, target=(world, stage)) + + +def main(): + level = os.environ.get("MARIO_LEVEL", "1-1") + env = make(level) + obs, done, info = env.reset(), False, {} + for line in sys.stdin: + request = json.loads(line) + if request["cmd"] == "reset": + if request.get("level", level) != level: + env.close() + level = request["level"] + env = make(level) + obs, done, info = env.reset(), False, {} + for _ in range(int(request.get("idle", 0))): # a different start: enemies in other phases + obs, _, done, info = env.step(0) + elif request["cmd"] == "step": + obs, _, done, info = env.step(int(request["buttons"])) + reply(env, obs, done, info) + + +if __name__ == "__main__": + main() diff --git a/Sources/JevMarioDemo/retro_bridge.py b/Sources/JevMarioDemo/retro_bridge.py new file mode 100644 index 0000000..723a828 --- /dev/null +++ b/Sources/JevMarioDemo/retro_bridge.py @@ -0,0 +1,157 @@ +"""Super Mario Bros. emulator bridge on stable-retro (unlimited save states, any level start), for JevMarioDemo. + +Requests, one JSON line each: + {"cmd": "reset", "level": "1-3", "idle": 0} start a level (optional idle frames) + {"cmd": "step", "buttons": 130} one frame + {"cmd": "run", "buttons": [130, 131, ...], "hold": true} a whole move in one call; with hold, the last input is + held while the game has control (pipe entry, ...) + {"cmd": "save"} / {"cmd": "load"} one save-state slot (the student's lookahead simulates + candidate moves from it, then restores); no image +Reply: one JSON line {"x", "screen_x", "screen_y", "on_ground", "done", "flag", "time", "frames", "images", "control", +"area", "transition"} then +`images` raw 240x256 RGB frames: one for reset / step, two for run (the last frame and the one 4 frames before it, +which the pixel student needs). Buttons are NES bitmasks (0x80 right, 0x40 left, 0x20 down, 0x10 up, 0x02 B, 0x01 A). + +Needs stable-retro and gym-super-mario-bros 7.4.0 (whose bundled ROM is copied in as retro's rom.nes; local use +only). demo.sh runs it through `uv run --with ...`. +""" +import collections +import json +import os +import shutil +import sys +import warnings + +import numpy as np + +warnings.filterwarnings("ignore") +import stable_retro as retro # noqa: E402 + +GAME = "SuperMarioBros-Nes-v0" +ORDER = ["B", None, "SELECT", "START", "UP", "DOWN", "LEFT", "RIGHT", "A"] +BITS = {"RIGHT": 0x80, "LEFT": 0x40, "DOWN": 0x20, "UP": 0x10, "START": 0x08, "SELECT": 0x04, "B": 0x02, "A": 0x01} + + +def install_rom(): + target = os.path.join(os.path.dirname(retro.data.__file__), "stable", GAME, "rom.nes") + if not os.path.exists(target): + import gym_super_mario_bros + + source = os.path.join(os.path.dirname(gym_super_mario_bros.__file__), "_roms", "super-mario-bros.nes") + shutil.copy(source, target) + + +def action(mask: int) -> np.ndarray: + return np.array([1 if name and mask & BITS[name] else 0 for name in ORDER], np.int8) + + +def area_for(world, stage): + return stage + 1 if world in (1, 2, 4, 7) and stage >= 2 else stage + + +class Game: + def __init__(self): + install_rom() + self.env = retro.make(GAME, state="Level1-1", use_restricted_actions=retro.Actions.ALL, render_mode=None) + self.env.reset() + self.u, self.em = self.env.unwrapped, self.env.unwrapped.em + + @property + def ram(self): + return self.u.get_ram() + + def poke(self, address, value): + self.u.data.memory.assign(address, "|u1", int(value)) + + def frame(self, screen=None) -> bytes: + out = np.zeros((240, 256, 3), np.uint8) + s = self.em.get_screen() if screen is None else screen + out[8:8 + s.shape[0], 8:8 + s.shape[1]] = s + return out.tobytes() + + def dying(self): + r = self.ram + return int(r[0x0E]) in (0x06, 0x0B) or int(r[0xB5]) > 1 + + def flag(self): + r = self.ram + return int(r[0x1D]) == 3 or int(r[0x770]) == 2 + + def reset(self, level, idle): + world, stage = (int(v) for v in level.split("-")) + self.env.reset() + for _ in range(30): + self.env.step(action(0)) + for address, value in ((0x75F, world - 1), (0x75C, stage - 1), (0x760, area_for(world, stage) - 1), + (0x75B, 0), (0x75A, 9), (0xB5, 2)): + self.poke(address, value) # target level, then a death reload starts it from the beginning + for _ in range(900): + self.env.step(action(0)) + r = self.ram + if int(r[0x770]) == 1 and int(r[0x0E]) == 8 and int(r[0x7A0]) == 0 and int(r[0xB5]) == 1: + break + for _ in range(idle): + self.env.step(action(0)) + + def step(self, mask): + self.env.step(action(mask)) + + def meta(self, frames, images, transition=False): + r = self.ram + done = self.flag() or self.dying() + return {"x": int(r[0x6D]) * 256 + int(r[0x86]), "screen_x": (int(r[0x86]) - int(r[0x71C])) % 256, + "screen_y": int(r[0xCE]), "on_ground": int(r[0x1D]) == 0, "done": done, "flag": self.flag(), + "time": int(r[0x7F8]) * 100 + int(r[0x7F9]) * 10 + int(r[0x7FA]), "frames": frames, "images": images, + "control": int(r[0x0E]) == 8, "area": int(r[0x750]) & 0x7F, "transition": transition} + + +def reply(meta, frames): + sys.stdout.buffer.write((json.dumps(meta) + "\n").encode()) + for f in frames: + sys.stdout.buffer.write(f) + sys.stdout.buffer.flush() + + +def main(): + game = Game() + game.reset(os.environ.get("MARIO_LEVEL", "1-1"), 0) + slot = game.em.get_state() + for line in sys.stdin: + request = json.loads(line) + cmd = request["cmd"] + if cmd == "reset": + game.reset(request.get("level", "1-1"), int(request.get("idle", 0))) + reply(game.meta(0, 1), [game.frame()]) + elif cmd == "step": + game.step(int(request["buttons"])) + reply(game.meta(1, 1), [game.frame()]) + elif cmd == "save": + slot = game.em.get_state() + reply(game.meta(0, 0), []) + elif cmd == "load": + game.em.set_state(slot) + reply(game.meta(0, 0), []) + elif cmd == "run": + masks, recent, n, transition = request["buttons"], collections.deque(maxlen=5), 0, False + for mask in masks: + game.step(int(mask)) + n += 1 + recent.append(game.em.get_screen().copy()) # padded only for the two frames sent back + if game.flag() or game.dying(): + break + if request.get("hold") and not (game.flag() or game.dying()): + for _ in range(300): # the game has control (pipe entry, area change): hold the last input + if int(game.ram[0x0E]) == 8: + break + transition = True + game.step(int(masks[-1])) + n += 1 + recent.append(game.em.get_screen().copy()) + if game.flag() or game.dying(): + break + last = game.frame(recent[-1]) if recent else game.frame() + reply(game.meta(n, 2, transition), [last, game.frame(recent[0]) if len(recent) == 5 else last]) + + +if __name__ == "__main__": + main()