Skip to content

About

GaeaPluginKit: MLXEngine packages as Gaea plugin workers (GPP v1), the manifest generator, and a tiny host check

Resources

Stars

0 stars

Watchers

0 watching

Forks

Latest commit

 

History

17 Commits

Folders and files

Repository files navigation

GaeaPluginKit

GaeaPluginKit turns MLXEngine model packages into a Gaea plugin worker. The worker is a signed executable in a .gaeaplugin bundle that speaks the Gaea Plugin Protocol v1 on stdin and stdout (xocialize-code/gaea, spec/GPP-1.md). The kit also writes the bundle's manifest.json from the plugin's own code, so the manifest can't drift from the worker.

Requires macOS 26 and mlx-engine-swift 0.72.0 or later (contract 1.57.0). Adapters need Gaea 0.7.0 or later.

Product What
GaeaPluginKit The worker: GPP's protocol loop over MLXServeEngine(workerBudget:), the plugin declaration, and the manifest generator
GaeaProtocol MLX-free: JSON values, the manifest's types, GPP's vocabulary. The Gaea host links no MLX and can share it
gaea-plugin-check A tiny MLX-free host. It verifies a built bundle, launches it under Process.launchRequirement, and drives a GPP session

A plugin is signed with its author's own Developer ID, and its manifest names that team. Gaea trusts its own team (GPP.trustedTeam) always, and any other once the operator approves it: from Gaea 0.13.0 a release from a new team is held for that answer. gaea-plugin-check checks the bundle against the manifest's team, and --team <ID> asserts which team that must be. The guide for plugin authors is in gaea-plugin-demo.

A plugin

import GaeaPluginKit

let plugin = GaeaPluginDefinition(
    id: "com.xocialize.gaea.z-image", version: "0.1.0", team: "CFZS6RXUCP",
    repo: "https://github.com/xocialize-code/gaea-plugin-z-image", title: "Z-Image", license: "Apache-2.0",
    executable: "gaea-z-image-worker",
    models: [.engine(ZImageTurboT2IPackage.self, id: "z-image-turbo", title: "Z-Image Turbo",
                     lanes: [LaneDefinition(id: "int4", default: true, configuration: ZImageConfiguration.turbo(quant: .int4))],
                     operations: [EngineOperation("image.generate", outputNames: ["image": "images"]) { job in
                         EngineCall(T2IRequest(prompt: try job.text("prompt") ?? "", seed: job.seed))
                     }])],
    offers: [/* PluginManifest.Offer(...) */])

await GaeaPluginMain.run(plugin)
  • <worker> serves GPP: what Gaea launches.

  • <worker> --manifest <file> writes manifest.json. models[] is MLXToolKit's own PackageManifest, plus one ResolvedLane per lane, as tools/manifest-gen writes them. The plugin build runs tools/validate.py --manifest on it before signing.

  • A text input (an MES prompt or instruction) arrives as a UTF-8 file: job.text("prompt") reads it (GPP-1 §3.3).

  • One engine call per result: EngineOperation(op, calls: { job in [EngineCall] }) runs them in order, for an MES count on a capability that returns one image per call. Outputs are listed in call order (image-0, image-1, …), effective values are merged, and each warning is kept once. Progress spans the whole job.

  • A self-running operation, EngineOperation(op) { job, session in … }, makes the engine calls it decides on with session.run(request) and returns its RunResult itself. session.run(request, model: id) calls another engine model of the same plugin, at its default lane, with its weights already on disk: the forge pack's guarded upscale runs NERVE under the words the text guard protects in a Nacre job.

  • Leave a lane configuration's models root unset. The engine fills it from the root Gaea sends in configure.

Audio: streams, voices and assets (0.7.0)

For MES Audio 0.1 (media-engine-spec v0.7.0, GPP-1 §3.3.1, §3.3.2, §4.3):

  • Streaming. When the host asks for an output as a stream (job.streams("audio")), send its samples as they come with context.emit("audio", samples:, sampleRate:, final:) (a custom model), or session.emit(…) in a perform operation, and return .streamed("audio"). Samples are floats in [-1, 1] at the model's own rate. Gaea converts them to the client's format. session.stream(ttsRequest) runs a StreamEmitting package (MLXEngine contract 1.25) and sends each of its chunks as it arrives.
  • Offers. Put the MES offer's supports and limits in its model ("supports": ["stream": true, "voices": 10, "sample_rate": 24000, "channels": 1, "seed": false], "limits": ["max_text_chars": 4000]), and the preset voices in Offer(…, voices: [...]) as MES Voice objects.
  • Assets. job.asset("voice_ref") is the MES asset id behind an input. Key anything cached from an input by it, and drop it in onAssetRemoved(asset, resident): resident.run(request, model:) sends a package its own eviction call while it is loaded, and does nothing when it isn't (0.8.0).

Video, previews and the weights' volume (0.9.0)

For Gaea's video phase (GPP-1 §3.2.1, §3.3, §3.3.3):

  • Video in and out. A video input arrives as an MP4 (H.264 or HEVC, with AAC audio where there is sound), whatever the client sent. Write outputs as MP4 or QuickTime files (.file("video", url, mediaType: "video/mp4")): Gaea reads their facts from the file. Engine packages take them as MLXToolKit's Video(format: .mp4, data:).
  • Previews. context.preview(fileURL) sends an image you've written under the run's output_dir (a frame, a draft) as the MES job's preview. Gaea reads the file and deletes it.
  • The weights' volume. An engine lane with a minSustainedReadBytesPerSecond floor is refused at prepare when the engine measures its volume below the floor. The kit sends that as storage_below_floor, with the measurement, and Gaea falls back to a lane the volume can feed. A prepared line carries the engine's measurement as storage.

Layer limits and supports (0.6.0)

An image.decompose offer (MES Image §3.2) says how many layers its model returns and whether it makes a composite: PluginManifest.Limits(maxLayerCount: 12) and Offer(…, supports: ["composite": true]), written as limits.max_layer_count and supports (GPP-1 §4.3). Absent supports keys are false.

Adapters (0.5.0)

A model whose package declares AdapterSupport on a surface takes LoRA adapters there (GPP-1 §3.3, §4.3, §4.8):

  • The manifest. Each offer gets adapter_support from the package's own AdapterSupport (base model, formats, kinds, max count, strength range), unless the offer sets one. A model can ship adapters from a registry in format 3 (the packages' registry/*.json): .engine(…).shipping(try AdapterRegistry.decode(data).adapters). Each is listed on every offer whose surface admits it, as a MES AdapterOffer in model.adapters (id <namespace>:<id>, title, kind, ic slots, strength, licence, the clearance note) with its adapters_meta (trigger, prompt template, status, clearance, format, pinned source). Compile the registry in with resources: [.embedInCode(…)]: the bundle carries no SwiftPM resource bundle.
  • A run. run.adapters arrive typed as job.adapters (JobAdapter: id, strength, kind, format, source, and an ic adapter's images by role). An operation puts try job.adapterRefs() on its T2IRequest or IEditRequest; it adds no triggers, which Gaea has already applied. Before the engine call the worker registers each adapter with its engine: a shipped one as its registry declares it, once Gaea's source matches the pin; an operator-registered one as a plain adapter on the surface's base model, under the licence Gaea already gated (the engine's enforcement is advisory). A missing file is fetched into <models_root>/adapters/<sha256>/<file> and reported as download_adapters progress; the engine checks its sha256. An operation that leaves the adapters off its request is refused rather than run without them.
  • Refusals. What the package's surface refuses (base model, kind, format, count, strength, conditioning) is invalid_request with the engine's reason; a digest mismatch is materialize_failed; an ic adapter the plugin doesn't ship is invalid_request, since GPP carries no slots or footprint.

gaea-plugin-demo is the template: bundle layout, signing, checks, notarization and release CI. gaea-plugin-z-image is the first plugin on an MLXEngine package.

gaea-plugin-check <bundle> [--operation op] [--model id] [--lane id] [--params json] [--model-params json] [--seed N] [--text name=text]… [--input name=path]… [--adapter id=strength]… [--adapter-input id:role=path]… [--models-root dir] [--budget bytes] [--timeout s] [--session file] checks a built bundle: its signature, a launch under the launch requirement, hello, and one run. --models-root points it at weights already on disk, and --budget must cover the lane. --adapter sends an adapter the offer ships as Gaea would, with the kind, format and source from its adapters_meta; --adapter-input gives an ic adapter an image for a slot.

What the worker does

GPP The kit
hello Plugin, models and engine from the definition. manifest_sha256 is the digest of the bundle's manifest.json
configure Builds MLXServeEngine(workerBudget:) and points its model store at models_root (adapters land under adapters/, fetched on demand). Settings get their manifest defaults filled in. The plugin's onConfigure hook then runs in a task of its own
materialize engine.materialize(package:), with materialize.progress from the download sink. Byte counts come from materializationPreview. A lane the package materializes itself is answered invalid_request: send prepare
prepare Registers the lane under the model's id, replacing another lane of the same model. Then engine.prepare; a download on the way is relayed from PreparationMonitor.events(). A custom model's lane is charged as an external tenant
run The run's adapters are registered with the engine and fetched (Adapters, above). The plugin's EngineOperation builds the canonical request, and engine.run executes it. A RunProgress sink feeds progress. The response is encoded with CapabilityWire.encoder(spool:) into output_dir: artifacts become files named after their field, other fields become values. Missing weights are not_materialized. A second run gets busy
cancel Cancels the run's task. The run ends with cancelled ("Cancelled at step k of n.")
unload engine.evict(package:), or the custom model's unload() and the withdrawal of its tenant
memory engine.workerMemoryReport(), unchanged: on memory.query, and unprompted after prepared, unloaded and a run's end
storage context.storage.get/put/delete/query: worker requests matched to the host's replies
shutdown, end of stdin Cancels a run in flight, then exits 0
anything else protocol errors for lines that aren't JSON or aren't a host message. invalid_request for a missing or unknown field, as the closed schema says

Per-job progress uses the RunProgress sink rather than RunMonitor.events(). The engine documents the sink as the per-job channel, and the monitor's stream is keyed by package. With one run per worker the two carry the same reports.

Tests

swift test

The tests need a checkout of xocialize-code/gaea, at ../gaea-mes-media-server or GAEA_SPEC_ROOT. They:

  • generate examples/v1/manifest/{demo,z-image}.json from code, check the result is the same JSON, and run validate.py --manifest on it;
  • replay every examples/v1/gpp/*.gpp: the H lines go to a worker over a real MLXServeEngine with fake packages, and each W line is compared field by field. Only pid, timings, measured memory and computed values are masked;
  • run each session the worker produced through validate.py --transcript;
  • cover the paths the examples don't walk: before configure, unknown fields, unknown lanes, unoffered operations, busy, and shutdown during a run;
  • cover adapters over a fake package that declares AdapterSupport: the registry decoder, the manifest's adapter_support, model.adapters and adapters_meta (validated), shipped and operator adapters fetched by digest and applied, an ic adapter's images in slot order, and every refusal.

Licence

Apache-2.0.

About

GaeaPluginKit: MLXEngine packages as Gaea plugin workers (GPP v1), the manifest generator, and a tiny host check

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages