# wd_tagger_mcp

一个专用于 WD Tagger 反推图片标签的 MCP 服务。支持通过 base64/data URL 或 HTTP/HTTPS URL 提交图片。端口、模型路径和标签路径都在 `config.toml` 中配置，其中文件路径必须使用相对路径。

## 目录

```text
.
  config.toml
  models/
    wd-eva02-large-tagger-v3/
      model.onnx
      selected_tags.csv
  wd_tagger_mcp/
    server.py
    tagger.py
    config.py
```

## 安装项目与依赖

建议在 Python 虚拟环境中安装和运行本项目。在项目根目录安装：

```text
python -m pip install -e .
```

## 模型下载

默认配置使用 Hugging Face 上的 [SmilingWolf/wd-eva02-large-tagger-v3](https://huggingface.co/SmilingWolf/wd-eva02-large-tagger-v3/tree/main)。只需下载以下两个文件：

- [model.onnx](https://huggingface.co/SmilingWolf/wd-eva02-large-tagger-v3/resolve/main/model.onnx?download=true)：ONNX 模型权重，约 1.26 GB。
- [selected_tags.csv](https://huggingface.co/SmilingWolf/wd-eva02-large-tagger-v3/resolve/main/selected_tags.csv?download=true)：与模型匹配的标签和分类信息。

下载后按以下目录结构放置：

```text
models/
  wd-eva02-large-tagger-v3/
    model.onnx
    selected_tags.csv
```

本服务通过 ONNX Runtime 载入 `model.onnx`，不需要下载仓库中的 `model.safetensors` 或 `model.msgpack`。如果将文件放在其他位置，请同步修改 `config.toml` 中的 `model_path` 和 `tags_path`；这两个配置项必须使用相对于 `config.toml` 的路径。

### 更换模型

本服务兼容 SmilingWolf 的以下官方模型。

v3：

- [wd-convnext-tagger-v3](https://huggingface.co/SmilingWolf/wd-convnext-tagger-v3/tree/main)
- [wd-vit-tagger-v3](https://huggingface.co/SmilingWolf/wd-vit-tagger-v3/tree/main)
- [wd-swinv2-tagger-v3](https://huggingface.co/SmilingWolf/wd-swinv2-tagger-v3/tree/main)
- [wd-vit-large-tagger-v3](https://huggingface.co/SmilingWolf/wd-vit-large-tagger-v3/tree/main)
- [wd-eva02-large-tagger-v3](https://huggingface.co/SmilingWolf/wd-eva02-large-tagger-v3/tree/main)

v2：

- [wd-v1-4-moat-tagger-v2](https://huggingface.co/SmilingWolf/wd-v1-4-moat-tagger-v2/tree/main)
- [wd-v1-4-swinv2-tagger-v2](https://huggingface.co/SmilingWolf/wd-v1-4-swinv2-tagger-v2/tree/main)
- [wd-v1-4-convnext-tagger-v2](https://huggingface.co/SmilingWolf/wd-v1-4-convnext-tagger-v2/tree/main)
- [wd-v1-4-convnextv2-tagger-v2](https://huggingface.co/SmilingWolf/wd-v1-4-convnextv2-tagger-v2/tree/main)
- [wd-v1-4-vit-tagger-v2](https://huggingface.co/SmilingWolf/wd-v1-4-vit-tagger-v2/tree/main)

切换模型时，下载对应仓库的 `model.onnx` 和 `selected_tags.csv`，放入独立目录，然后修改 `config.toml`：

```toml
[model]
model_path = "models/wd-convnext-tagger-v3/model.onnx"
tags_path = "models/wd-convnext-tagger-v3/selected_tags.csv"
```

每个模型都必须使用同一个仓库中与其配套的 `selected_tags.csv`，不能在 v2、v3 或不同模型之间混用模型和标签文件。服务会自动读取 ONNX 模型的输入尺寸，并在启动时校验输入布局、输出标签数与 CSV 是否匹配。修改配置后需要重启服务。

## 启动

在项目根目录启动 MCP 服务：

```text
python -m wd_tagger_mcp.server --config config.toml
```

该命令会在当前终端中前台运行服务。当前 `config.toml` 使用 SSE transport，监听地址为：

```text
0.0.0.0:48765
```

服务日志写入 `config.toml` 中 `[logging].log_path` 指定的文件，默认为：

```text
logs/wd-tagger-mcp-server.log
```

要停止服务，在运行它的终端中按 `Ctrl+C`。

默认保留 Danbooru 标签中的下划线，例如 `looking_at_viewer`。如果需要输出更接近自然语言的 `looking at viewer`，可在 `config.toml` 中设置：

```toml
[tagging]
replace_underscore = true
```

## 工具

- `wd_tagger_tag_image_base64`: 传入 base64 或 data URL，返回 prompt、general tags、character tags、rating tags。
- `wd_tagger_tag_image_url`: 传入 HTTP/HTTPS 图片 URL，服务端下载图片后返回 prompt、general tags、character tags、rating tags。

示例参数：

```json
{
  "image_base64": "data:image/png;base64,iVBORw0KGgo...",
  "general_threshold": 0.35,
  "character_threshold": 0.85,
  "include_rating": false
}
```

```json
{
  "image_url": "https://example.com/image.png",
  "general_threshold": 0.35,
  "character_threshold": 0.85,
  "include_rating": false
}
```