spb/zyquo-mlx Public MIT

1//2//  ModelStore.swift3//  Zyquo MLX4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//89import Foundation1011/// Errors surfaced when a directory is not a usable MLX model.12enum ModelValidationError: LocalizedError {13    case notADirectory(URL)14    case missingConfig(URL)15    case missingWeights(URL)16    case inconsistentShards(missing: [String])1718    var errorDescription: String? {19        switch self {20        case .notADirectory(let url):21            "Not a folder: \(url.path)"22        case .missingConfig(let url):23            "No config.json found in \(url.lastPathComponent) — this is not an MLX model folder."24        case .missingWeights(let url):25            "No weight files (*.safetensors) found in \(url.lastPathComponent)."26        case .inconsistentShards(let missing):27            "The weight index references missing shards: \(missing.joined(separator: ", "))."28        }29    }30}3132/// Scans, validates, and describes the local model library33/// (`~/Library/Application Support/ZyquoMLX/Models/`).34///35/// Validation follows the on-disk MLX format (docs/MLX-RESEARCH.md §5.3):36/// `config.json` (+ optional `quantization` dict), `model*.safetensors`37/// (single or sharded + index), tokenizer files.38actor ModelStore {3940    static let shared = ModelStore()4142    private let root: URL4344    init(root: URL = PersistenceService.modelsDirectory) {45        self.root = root46    }4748    /// All valid models in the library. Invalid directories are skipped49    /// (they surface through `validate` when the user targets them directly).50    func scan() throws -> [LocalModel] {51        let fm = FileManager.default52        guard fm.fileExists(atPath: root.path) else { return [] }53        let entries = try fm.contentsOfDirectory(54            at: root, includingPropertiesForKeys: [.isDirectoryKey],55            options: .skipsHiddenFiles)56        return entries.compactMap { try? describe(directory: $0) }57            .sorted { $0.installedAt > $1.installedAt }58    }5960    /// Validate and describe one model directory (works for any path, not61    /// just the library root — Playground can open arbitrary folders).62    func describe(directory: URL) throws -> LocalModel {63        let fm = FileManager.default64        var isDir: ObjCBool = false65        guard fm.fileExists(atPath: directory.path, isDirectory: &isDir), isDir.boolValue else {66            throw ModelValidationError.notADirectory(directory)67        }6869        let configURL = directory.appendingPathComponent("config.json")70        var config: [String: Any]?71        if let data = try? Data(contentsOf: configURL) {72            config = try JSONSerialization.jsonObject(with: data) as? [String: Any]73        }7475        let type = ModelType.detect(directory: directory, config: config)7677        // Image-generation layouts (FLUX) keep weights in component subdirs and78        // have no root config; everything else must have config.json.79        if config == nil && type != .imageGeneration {80            throw ModelValidationError.missingConfig(directory)81        }8283        let allFiles = (try? fm.subpathsOfDirectory(atPath: directory.path)) ?? []84        let weightFiles = allFiles.filter { $0.hasSuffix(".safetensors") || $0.hasSuffix(".npz") }85        guard !weightFiles.isEmpty else {86            throw ModelValidationError.missingWeights(directory)87        }8889        // Shard-index consistency (model.safetensors.index.json → weight_map).90        let indexURL = directory.appendingPathComponent("model.safetensors.index.json")91        var totalParameters: Int64?92        if let indexData = try? Data(contentsOf: indexURL),93            let index = try? JSONSerialization.jsonObject(with: indexData) as? [String: Any]94        {95            if let weightMap = index["weight_map"] as? [String: String] {96                let referenced = Set(weightMap.values)97                let present = Set(weightFiles.map { ($0 as NSString).lastPathComponent })98                let missing = referenced.subtracting(present).sorted()99                // Some mlx-community repos ship a consolidated100                // model.safetensors next to a stale sharded index (seen live:101                // Qwen3-VL-4B-Instruct-4bit) — treat that as valid.102                if !missing.isEmpty && !present.contains("model.safetensors") {103                    throw ModelValidationError.inconsistentShards(missing: missing)104                }105            }106            if let meta = index["metadata"] as? [String: Any],107                let params = meta["total_parameters"] as? Int64 ?? (meta["total_parameters"] as? Int).map(Int64.init)108            {109                totalParameters = params110            }111        }112113        var weightsSize: Int64 = 0114        var diskSize: Int64 = 0115        for path in allFiles {116            let fileURL = directory.appendingPathComponent(path)117            guard let size = (try? fileURL.resourceValues(forKeys: [.fileSizeKey]))?.fileSize118            else { continue }119            diskSize += Int64(size)120            if path.hasSuffix(".safetensors") || path.hasSuffix(".npz") {121                weightsSize += Int64(size)122            }123        }124125        var quantization: QuantizationInfo?126        if let quantDict = config?["quantization"] as? [String: Any],127            let data = try? JSONSerialization.data(withJSONObject: quantDict)128        {129            quantization = try? JSONDecoder().decode(QuantizationInfo.self, from: data)130        }131132        let name = directory.lastPathComponent133        let installedAt =134            (try? directory.resourceValues(forKeys: [.creationDateKey]))?.creationDate ?? .now135136        return LocalModel(137            id: name,138            name: name,139            repoID: Self.repoID(fromDirectoryName: name),140            type: type,141            directory: directory,142            architecture: config?["model_type"] as? String,143            parameterCount: totalParameters ?? Self.parameterCount(fromName: name, quantization: quantization, weightsSize: weightsSize),144            quantization: quantization,145            weightsSize: weightsSize,146            diskSize: diskSize,147            installedAt: installedAt148        )149    }150151    /// Delete a model from the library.152    func delete(_ model: LocalModel) throws {153        try FileManager.default.removeItem(at: model.directory)154    }155156    // MARK: - Name parsing157158    /// Library directories are named "org--repo" (HF convention for local caches).159    private static func repoID(fromDirectoryName name: String) -> String? {160        guard name.contains("--") else { return nil }161        return name.replacingOccurrences(of: "--", with: "/")162    }163164    /// Parse "…-4B-…" / "…-0.6B-…" / "…-135M-…" out of a repo name; fall back165    /// to estimating from weight bytes (docs/MODELS.md §3 bytes-per-param).166    static func parameterCount(fromName name: String, quantization: QuantizationInfo?, weightsSize: Int64) -> Int64? {167        if let range = name.range(of: #"(\d+(?:\.\d+)?)\s*[bB](?=[-_.]|$)"#, options: .regularExpression) {168            let token = name[range].dropLast()169            if let value = Double(token.trimmingCharacters(in: .whitespaces)) {170                return Int64(value * 1_000_000_000)171            }172        }173        if let range = name.range(of: #"(\d+(?:\.\d+)?)\s*[mM](?=[-_.]|$)"#, options: .regularExpression) {174            let token = name[range].dropLast()175            if let value = Double(token.trimmingCharacters(in: .whitespaces)) {176                return Int64(value * 1_000_000)177            }178        }179        guard weightsSize > 0 else { return nil }180        let bytesPerParam = MemoryAdvisor.bytesPerParameter(quantization: quantization)181        return Int64(Double(weightsSize) / bytesPerParam)182    }183}184