ai-toolkit-mcp

An MCP server for ostris/ai-toolkit. It lets an AI agent in Griptape Nodes, Claude Code, Claude Desktop, or any other MCP client set up, launch, and monitor LoRA training runs on your own machine.

"Train a LoRA of my dog on the rufus dataset, same settings as last week's run but 2000 steps. Check on it in a bit and show me the latest samples."

It talks to the ai-toolkit web UI's API, so every job it creates shows up in the UI, uses the UI's per-GPU queue, and appears on the UI's loss graph. You can switch between chatting and clicking at any time.

Requirements

• ai-toolkit installed, with its UI running (cd ai-toolkit/ui && npm run build_and_start, default address http://localhost:8675)
• uv

You don't need to clone this repo. uvx fetches and runs it.

Setup

Griptape Nodes

1. Open Settings → MCP Servers → + New MCP Server

2. Server Name/ID: ai-toolkit, Connection Type: Local Process (stdio)

3. Configuration JSON:

   {
     "transport": "stdio",
     "command": "uvx",
     "args": ["--from", "git+https://github.com/shhlife/ai-toolkit-mcp", "ai-toolkit-mcp"],
     "env": { "AI_TOOLKIT_URL": "http://localhost:8675" }
   }
   

4. Enable the ai-toolkit server on an Agent node (or use it from an MCP Task node) and start asking.

This lets you build workflows like: generate or caption a dataset with Griptape nodes, have the Agent kick off training, then pull the finished LoRA back into a generation flow.

Claude Code

claude mcp add ai-toolkit -- uvx --from git+https://github.com/shhlife/ai-toolkit-mcp ai-toolkit-mcp

Claude Desktop (or any MCP client using mcpServers JSON)

{
  "mcpServers": {
    "ai-toolkit": {
      "command": "uvx",
      "args": ["--from", "git+https://github.com/shhlife/ai-toolkit-mcp", "ai-toolkit-mcp"]
    }
  }
}

Configuration

Tools

Read-only tools are marked as read-only, so clients with tool approval (Griptape Nodes, Claude) can auto-approve them and only ask before anything that creates, starts, or stops a job. Nothing in this server deletes jobs, datasets, or files (set_captions does overwrite the captions you give it).

create_dataset uploads through the UI's API rather than writing into ai-toolkit's folder directly, so it works the same whether ai-toolkit is on your machine or a remote GPU box. ai-toolkit replaces spaces and special characters in uploaded file names with _.

How agents learn best practices

• training_guide returns practical advice per topic. The tool descriptions tell agents to call it before planning a run, so this works in every MCP client. The text lives in guide.py; edit it as you learn what works.
• Model presets come from your ai-toolkit install's own model definitions, the same ones the UI's model dropdown uses, so new models show up automatically when you update ai-toolkit.
• create_job / update_job check the config and return warnings in plain language.
• The server also sends usage instructions when a client connects. Claude Code reads those; Griptape Nodes currently doesn't pass them to its agents, which is why the guide is a tool.

Reusing models you already downloaded

ai-toolkit keeps ComfyUI-format weights in its own models folder for many models and downloads them there even when the model is already in your Hugging Face cache (LTX-2.5 always loads this way). When ai-toolkit runs on the same machine as this server, create_job / update_job point the job at the cached files instead and say so under using_hf_cache. Models that aren't cached get a warning before anything starts. The cache location comes from HF_HUB_CACHE, HUGGINGFACE_HUB_CACHE, or HF_HOME.

Tips

• create_job(model=...) applies the same settings the UI does when you pick a model. For edit models that need control images (Kontext, Qwen-Image-Edit), set one job up in the UI and have the agent clone it with base_job.
• Per-model reference configs live in ai-toolkit's config/examples/ folder. You can paste one into create_job(config=...).

Development

uv sync
uv run pytest
uv run mcp dev src/ai_toolkit_mcp/server.py   # MCP Inspector

License

MIT