[Tagger] 数据集打标 — 一起打磨默认模型 / 下载进度 / API 打标

[wochenlong]

背景

想和 @niangao2331 一起把 WebUI 「数据集打标」（侧栏 Tagger 页，批量 WD/CL 打标）这块做得更顺手。你之前修的 #23（整合包 Download-Anima-Model.bat 路径套娃）已经帮了大忙，如果你愿意继续盯这块，我们就以这个 Issue 当协作白板：先对齐目标，再各自 PR 合进来。

有想法、想改范围，直接在下面评论就好。

---

怎么协作（团队内约定）

1. 开发方式：fork → 自己的分支（例如 feature/tagger）→ PR 到 main；main 我们这边 merge，避免多人直接推。
2. 范围说明：这个 Issue 只管 「数据集打标」，侧栏 「标签编辑」（mikazuki/dataset-tag-editor/，Gradio 那块）是另一条线，暂时别绑在一起，免得互相踩脚。
3. 目录契约：mikazuki/tagger/、API、schema 都可以大胆加；根目录启动脚本、python_embeded/、SD-Trainer/ 这类路径见 docs/repo-layout.md，动之前咱们说一声就行。
4. 发挥空间：Phase 1 下面的条目是「当前最想做的」，不是死规定；UX、模型列表、文档、测试你看着办，只要不拆训练核心、不动上面说的契约路径，后面 Phase 2 也可以你自己拆 Issue。

---

眼下想一起做的（Phase 1）

1. 默认模型选型和固化

• Danbooru 向：一起看下现在哪套 WD 当默认最合适（现状是 wd14-convnextv2-v2 / SmilingWolf/wd-v1-4-convnextv2-tagger-v2，你有更好选型可以直接换）。
• 自然语言向：希望至少接一种 caption/NL 模型（BLIP、Florence 或其他 HF 上的都行），和 Danbooru 标签在 UI 和输出上区分清楚。
• 动默认模型时，顺带对齐：mikazuki/tagger/defaults.py、schema 默认项、scripts/prefetch_default_tagger.py、build-scripts/build_portable.ps1 预下载列表。

2. 模型下载进度

小白第一次下 HF 模型时经常不知道在干嘛。目标：WebUI 或控制台能看到进度（字节或分文件阶段都行）。
实现上你定：huggingface_hub 回调、SSE/WebSocket、或给 mikazuki/app/api.py 加个轻量状态接口都行，别卡住训练那条 API 就好。

3. API 打标（例如 Gemini）

• 和本地 ONNX/WD 并列的云端路径，首期 Gemini 多模态就挺好。
• Key 建议走环境变量或本地 config/（别进 git）；超时、限速、失败重试、费用提示可以按你习惯来。
• 输出继续写图片旁的 .txt，和现有训练数据集习惯一致。

---

以后有空可以再聊（Phase 2+）

• 任务队列 / 暂停取消 / 已打标跳过策略
• 整合包体积 vs 离线模型包策略
• 打标结果预览、抽样质检
• 和「标签编辑」页只读衔接（不耦合 Gradio）

---

代码地图：这块 mainly 动这里 ✅

区域	路径	说明
打标核心	mikazuki/tagger/	interrogator.py、interrogators/*、format.py、dbimutils.py、defaults.py
预下载	scripts/prefetch_default_tagger.py	安装/启动/构建时预拉默认模型
API	mikazuki/app/api.py	POST /interrogate 及进度/状态类接口
请求模型	mikazuki/app/models.py	TaggerInterrogateRequest 扩展
表单 Schema	mikazuki/schema/tagger.ts	下拉、默认值、新开关
打标页文案	frontend/dist/tagger.html、frontend/dist/assets/tagger.html.*.js	说明区、进度 UI（patch 约定见 frontend/VENDOR.md）
文档	README-zh.md 打标 FAQ	用户可见说明
整合包构建	build-scripts/build_portable.ps1	内置模型、[3/6] 预下载
安装/启动	install-cn.ps1、install.ps1、run_gui_source.bat、scripts/portable/launch_portable.bat	预取逻辑可扩展

入口：POST /interrogate → on_interrogate()（mikazuki/tagger/interrogator.py），后台跑，标签写同目录 .txt。

缓存：HF_HOME=huggingface → huggingface/hub/（和 wd14.py 里 hf_hub_download 一致）。

---

尽量少碰 / 动之前喊一声 ⚠️

路径	说明
mikazuki/dataset-tag-editor/	标签编辑，另一条产品线
vendor/sd-scripts/ 等训练	训练引擎
gui.py、train_monitor/	全局主流程，除非真要挂服务
run_gui.bat、python_embeded/、SD-Trainer/ 目录名	整合包契约
start_autodl.sh（根路径）	云镜像写死
legacy/cli/tagger.ps1 等	旧 CLI，参考即可

frontend/dist/ 没有上游源码构建权时，可以 patch 预编译 JS/HTML（记得记 hash，见 frontend/VENDOR.md）；要大改 UI 咱们先对一下。

---

当前 main 上的基线（方便你 fork）

• 默认 Danbooru：wd14-convnextv2-v2（revision=v2.0）
• 源码：install-cn.ps1 + run_gui.bat 会跑 prefetch_default_tagger.py --if-missing
• 整合包构建会打进 huggingface/hub/ 默认模型

相关：6130033 及之后 CONTRIBUTORS 更新。

---

Phase 1 做完时咱们怎么算「可以 merge」

• 默认 Danbooru 模型对新用户够用，README/页内说明跟得上
• 至少一种 NL/caption 模型可选，输出格式写清楚
• 大模型下载有可见进度（UI 或约定好的日志二选一+）
• API 打标（如 Gemini）配好 Key 后能跑通一个目录
• PR 里简单写自测环境、是否影响整合包体积，有截图/日志更好
• 没动 dataset-tag-editor、run_gui.bat 契约还在

---

@niangao2331 有兴趣的话直接在这个 Issue 里回复分工或开分支；PR 里写 Fixes #40 会自动关联。有卡点 @wochenlong 或群里喊一声就行。

[wochenlong]

负责人对齐（2026-05-24）

@niangao2331 这期 只做 backend，不写 frontend/dist、不改侧栏。UI 与「标签编辑 / 数据集合一」由 @wochenlong 另线（#41）。

---

产品背景（不阻塞本 Issue）

侧栏中长期会收成 「数据集」 一项，页内 Tab：自动打标 | 标签编辑。

• [Tagger] 数据集打标 — 一起打磨默认模型 / 下载进度 / API 打标 #40 → 自动打标 backend（UI 暂用现有 /tagger.md）
• 阿里云镜像下载pytorch异常慢 #41 → 共享 scan/caption API + 编辑 UI + 侧栏合并

---

功能范围（Must Have）

1. WD 本地打标（延续 + 加固）

• 默认模型：维持或优化 wd14-convnextv2-v2（mikazuki/tagger/defaults.py）
• 输出：图片同目录 .txt sidecar
• 冲突策略：ignore / copy / prepend，行为可测
• 沿用现有 TaggerInterrogateRequest 字段，向后兼容

2. 请求模型扩展（Tag / NL / API 统一入口）

在不破坏现有调用前提下扩展 TaggerInterrogateRequest（或 union）：

字段	类型	说明
mode	tag | natural_language	tag 列表 vs 自然语言 caption
engine	local_wd | local_nl | api_gemini	执行引擎
scope	all | missing_caption | selected	selected 需 image_paths[]
image_paths	string[]?	scope=selected 时必填
prompt_preset	string?	NL/API：generic / anima_detailed / short 等
caption_extension	string	默认 .txt

Phase 1 最低：mode=tag+engine=local_wd（等价现状）；另至少一种 natural_language 或 api_gemini。

3. 自然语言本地打标（NL）

• 至少一种 HF 本地 caption 模型（BLIP / Florence / ONNX 友好，你选型）
• 输出写 .txt；README 说明与 Danbooru tag 的区别

4. API 打标（Gemini）

• Google Gemini Vision；Key 走 env 或 config/（禁止进 git）
• 超时、重试、429 退避；单图失败不拖死整批
• README-zh：配置与费用提示

5. 模型下载进度

• GET /api/tagger/download-status（或扩展 status 的 download 字段）
• HF 首次下载可见文件/字节进度；不阻塞训练 API

6. 打标任务进度

• GET /api/tagger/status：phase = idle / downloading / tagging / done / error
• Gemini/NL 路径也更新 tagger_progress.py

---

明确不做（Out of Scope）

• frontend/dist/*、scripts/patch-*.py
• mikazuki/dataset-tag-editor/、:28001 Gradio
• 标签编辑、批量 tag 替换/排序/去重（阿里云镜像下载pytorch异常慢 #41）
• GET /api/dataset/scan 图库索引（阿里云镜像下载pytorch异常慢 #41）
• 侧栏「数据集」合一壳子（阿里云镜像下载pytorch异常慢 #41）
• run_gui.bat、python_embeded/、SD-Trainer/ 整合包契约路径
• vendor/sd-scripts/ 训练引擎

---

可动代码范围

区域	路径
打标核心	mikazuki/tagger/
NL/API 引擎	mikazuki/tagger/interrogators/（新增）
API	mikazuki/app/api.py（改前 comment 一声）
请求模型	mikazuki/app/models.py
Schema	mikazuki/schema/tagger.ts（仅定义，可不 patch dist）
预下载	scripts/prefetch_default_tagger.py
整合包	build-scripts/build_portable.ps1
测试	tests/test_tagger*.py 必须新增
API 文档	docs/api/dataset-tagging.md 必须新增
用户文档	README-zh.md 打标 FAQ

---

交付物（PR 必含）

代码

• API 扩展 + WD 兼容
• NL 或 Gemini 至少一种跑通
• download-status + progress 覆盖全流程

文档

• docs/api/dataset-tagging.md（请求/响应 JSON 示例）
• README-zh Gemini Key 说明

测试

• tests/test_tagger_api.py：冲突策略、写 .txt、mock Gemini

PR 描述：Fixes #40、curl 自测、pytest 结果、是否影响整合包体积

---

验收标准（merge 前）

1. 不传 mode/engine 时与现网 main 行为一致
2. 负责人 curl 可测 WD + 至少一种 NL/API
3. 下载/打标进度可见
4. pytest tests/test_tagger_api.py -q 通过
5. 禁区未动

---

建议 PR 拆分

PR	内容
PR-1	mode/engine/scope 扩展 + WD 兼容 + 测试基架
PR-2	下载进度 + prefetch
PR-3	local_nl 或 api_gemini（二选一）
PR-4	另一种引擎 + README

分支：feature/tagger-#40-<短描述>

---

自测命令

python gui.py --listen

curl -X POST http://127.0.0.1:28000/api/interrogate `
  -H "Content-Type: application/json" `
  -d '{"path":"D:/path/to/dataset","interrogator_model":"wd14-convnextv2-v2","batch_input_recursive":true,"batch_output_action_on_conflict":"copy"}'

curl http://127.0.0.1:28000/api/tagger/status
curl http://127.0.0.1:28000/api/tagger/download-status

python -m pytest tests/test_tagger_api.py -q

---

请回复确认

• 认领 [Tagger] 数据集打标 — 一起打磨默认模型 / 下载进度 / API 打标 #40 全部 backend
• 第一 PR 做：___
• 预计时间：___
• NL 模型选型：___
• Gemini Key 存储方案：___