spb/zyquo-mlx Public MIT

1//2//  ConversionService.swift3//  Zyquo MLX4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//89import Foundation10import MLX11import MLXLLM12import MLXLMCommon1314/// Progress for conversion/quantization/fusion jobs.15enum ConversionEvent: Sendable {16    case stage(String, fraction: Double?)17    case finished(outputDirectory: URL)18    case failed(message: String)19}2021enum ConversionServiceError: LocalizedError {22    case alreadyQuantized23    case outputExists(URL)2425    var errorDescription: String? {26        switch self {27        case .alreadyQuantized:28            "This model is already quantized — de-quantize first if you want a different recipe."29        case .outputExists(let url):30            "Output folder already exists: \(url.lastPathComponent). Choose another name or delete it."31        }32    }33}3435/// Conversion, quantization, and adapter fusion (charter 3.C).36///37/// Execution split per docs/MLX-RESEARCH.md §2/§6:38/// - Quantizing a local safetensors LLM → **native Swift**39///   (`LLMModelFactory.convert`, affine/mxfp4).40/// - Fusing adapters to disk, HF→MLX conversion of arbitrary repos41///   (incl. `.bin`), de-quantize → **Python bridge** (`mlx_lm` pinned).42actor ConversionService {4344    static let shared = ConversionService()4546    // MARK: - Swift-native quantization4748    /// Quantize a local, unquantized safetensors LLM into a new model49    /// directory in the library.50    func quantize(model: LocalModel, config: QuantConfig, outputName: String)51        -> AsyncStream<ConversionEvent>52    {53        let output = PersistenceService.modelsDirectory54            .appendingPathComponent(outputName, isDirectory: true)5556        let (stream, continuation) = AsyncStream.makeStream(of: ConversionEvent.self)57        Task {58            do {59                guard model.quantization == nil else { throw ConversionServiceError.alreadyQuantized }60                guard !FileManager.default.fileExists(atPath: output.path) else {61                    throw ConversionServiceError.outputExists(output)62                }63                let options = ModelConversionOptions(64                    bits: config.bits,65                    groupSize: config.groupSize,66                    mode: config.mode == .mxfp4 ? .mxfp4 : .affine67                )68                _ = try await LLMModelFactory.shared.convert(69                    from: model.directory,70                    to: output,71                    options: options,72                    progressHandler: { progress in73                        continuation.yield(74                            .stage(progress.stage.rawValue, fraction: progress.fractionCompleted))75                    }76                )77                Memory.clearCache()78                continuation.yield(.finished(outputDirectory: output))79            } catch {80                continuation.yield(.failed(message: error.localizedDescription))81            }82            continuation.finish()83        }84        return stream85    }8687    // MARK: - Python-bridge pipelines8889    /// Fuse a trained adapter into its base model, producing a standalone90    /// MLX model directory in the library.91    func fuse(baseModel: LocalModel, adapterDirectory: URL, outputName: String, dequantize: Bool)92        async throws -> AsyncStream<ConversionEvent>93    {94        let output = PersistenceService.modelsDirectory95            .appendingPathComponent(outputName, isDirectory: true)96        guard !FileManager.default.fileExists(atPath: output.path) else {97            throw ConversionServiceError.outputExists(output)98        }99100        var arguments = [101            "--model", baseModel.directory.path,102            "--adapter-path", adapterDirectory.path,103            "--save-path", output.path,104        ]105        if dequantize { arguments.append("--dequantize") }106107        return try await pythonJob(script: "zyquo_fuse", arguments: arguments, output: output)108    }109110    /// Convert an HF repo/directory (any weights format mlx-lm handles) to111    /// MLX, optionally quantized.112    func convert(hfPath: String, outputName: String, quantize: QuantConfig?)113        async throws -> AsyncStream<ConversionEvent>114    {115        let output = PersistenceService.modelsDirectory116            .appendingPathComponent(outputName, isDirectory: true)117        guard !FileManager.default.fileExists(atPath: output.path) else {118            throw ConversionServiceError.outputExists(output)119        }120121        var arguments = ["--hf-path", hfPath, "--mlx-path", output.path]122        if let quantize {123            arguments += [124                "--quantize",125                "--q-bits", String(quantize.bits),126                "--q-group-size", String(quantize.groupSize),127                "--q-mode", quantize.mode.rawValue,128            ]129        }130131        return try await pythonJob(script: "zyquo_convert", arguments: arguments, output: output)132    }133134    private func pythonJob(script: String, arguments: [String], output: URL)135        async throws -> AsyncStream<ConversionEvent>136    {137        let pythonStream = try await PythonRunner.shared.stream(138            script: script, arguments: arguments)139        let (stream, continuation) = AsyncStream.makeStream(of: ConversionEvent.self)140        Task {141            do {142                for try await event in pythonStream {143                    switch event.event {144                    case "start":145                        continuation.yield(.stage(event.string("stage") ?? "working", fraction: nil))146                    case "done":147                        continuation.yield(.finished(outputDirectory: output))148                    case "error":149                        continuation.yield(.failed(message: event.string("message") ?? "unknown"))150                    default:151                        break152                    }153                }154            } catch {155                continuation.yield(.failed(message: error.localizedDescription))156            }157            continuation.finish()158        }159        return stream160    }161}162