xsl
ce445b64a6
部署修复： - torch.load 增加 weights_only=False patch，兼容 PyTorch 2.6+ 加载旧权重 - OSS 改为懒加载，本地用 output_format=base64 无需配凭证即可启动 - 补全被 gitignore 误排除的必需代码：core/models/layers/data、models/layers/data、keypoints/lib - webui 训练命令 --xformers 改 --sdpa（修复 xformers 无 CUDA 支持报错） 功能调整： - hair_grow_service 端口改 8899、preview 路由修复（send_file） - list_hairstyles 增加发型白名单，测试页只展示当前5个发型 新增脚本： - train_lora_parallel.py：直接调 kohya 并行训练 LoRA（绕过 photo_service 串行限制） - train_hairstyles_parallel.py / train_batch_stepC.py：批量训练辅助脚本 - scripts/sync_data_to_server.sh：大文件断点续传到云服务器 文档： - docs/换发型集成文档.md：换发型完整流程、服务架构、资源依赖、训练方法、集成步骤
409 lines
18 KiB
Markdown
409 lines
18 KiB
Markdown
# 换发型功能集成文档
|
||
|
||
本文档说明"换发型"功能的完整工作流程、服务依赖、资源结构，供其他项目集成参考。
|
||
|
||
---
|
||
|
||
## 一、功能概述
|
||
|
||
**换发型**：用户上传一张人像照片，选择一个目标发型（hair_id），系统把用户的头发替换成目标发型，保留人脸和身体。
|
||
|
||
- **输入**：用户人像图（base64 或 URL） + 目标发型 ID
|
||
- **输出**：换好发型的结果图（base64 或 OSS URL）
|
||
- **耗时**：单次约 10-12 秒（GPU: RTX 5090）
|
||
|
||
### 同类功能（同套服务支撑）
|
||
- **换发色**（`/hairColor/v2`）：把头发颜色换成指定 RGB，额外依赖 `ref_color` 颜色查找表
|
||
- **训练新发型**：用发型图片训练新的 LoRA，详见第七节
|
||
|
||
---
|
||
|
||
## 二、服务架构
|
||
|
||
换发型由 **3 个后端服务**协同完成（缺一不可），另有 1 个前端页面服务（可选）。
|
||
|
||
| 服务 | 端口 | conda 环境 | 职责 |
|
||
|------|------|-----------|------|
|
||
| **webui** | 57860 | `sdwebui` | Stable Diffusion img2img 推理引擎（基于 AUTOMATIC1111 webui） |
|
||
| **photo_service** | 32678 | `py310` | LoRA 调度：推理时把对应发型的 LoRA 复制到 webui 并注入 prompt；训练时调度 kohya |
|
||
| **hair_service_sd** | 8801 | `my_hair` | 换发型主服务：人脸检测、关键点、头发抠图、材质推理、图像合成 |
|
||
| hair_grow_service | 8899 | `my_hair` | （可选）前端测试页面，提供发型列表/预览/换发型 UI |
|
||
|
||
### 服务间调用链（换发型）
|
||
|
||
```
|
||
调用方
|
||
│ POST /api/swapHair/v1
|
||
▼
|
||
hair_service_sd (8801) ← 人脸检测/关键点/抠图/材质推理/图像合成
|
||
│ POST /api/hair/inference (转发 img2img 请求)
|
||
▼
|
||
photo_service (32678) ← 复制 LoRA 到 webui、prompt 注入 <lora:xxx:1.0>
|
||
│ POST /sdapi/v1/img2img
|
||
▼
|
||
webui (57860) ← SD img2img 真正推理
|
||
```
|
||
|
||
### 启动方式
|
||
|
||
```bash
|
||
cd /home/xsl/change_hair
|
||
bash start_all.sh start # 启动 webui + photo_service + hair_service_sd
|
||
# 等待 1-2 分钟（webui 加载 SD 模型）
|
||
|
||
# 可选：前端页面
|
||
nohup bash start_hairgrow.sh > project/logs/hair_grow_service.log 2>&1 &
|
||
```
|
||
|
||
> 三个核心服务有启动顺序依赖：webui 必须先启动并加载完模型（约 1-2 分钟），photo_service 和 hair_service_sd 才能正常工作。
|
||
|
||
---
|
||
|
||
## 三、换发型 API 接口
|
||
|
||
### 请求
|
||
|
||
```
|
||
POST http://<host>:8801/api/swapHair/v1
|
||
Content-Type: application/json
|
||
```
|
||
|
||
**请求体参数：**
|
||
|
||
| 参数 | 类型 | 必填 | 说明 |
|
||
|------|------|:---:|------|
|
||
| `hair_id` | string | ✓ | 目标发型 ID（如 `chang_tuoyuan`），对应发型资源目录名 |
|
||
| `task_id` | string | ✓ | 任务唯一标识，用于临时文件命名 |
|
||
| `user_img_path` | string | ✓ | 用户人像图。支持 `data:image/jpeg;base64,<base64>` 或图片 URL |
|
||
| `is_hr` | string | ✓ | 是否高清模式。`"true"` 输出 1152×1536，`"false"` 输出 576×768 |
|
||
| `output_format` | string | ✗ | 输出格式。`"base64"`（推荐，直接返回图）或 `"url"`（上传 OSS 返回链接，需配 OSS 凭证） |
|
||
|
||
**请求示例：**
|
||
|
||
```bash
|
||
curl -X POST http://127.0.0.1:8801/api/swapHair/v1 \
|
||
-H "Content-Type: application/json" \
|
||
-d '{
|
||
"hair_id": "chang_tuoyuan",
|
||
"task_id": "req_001",
|
||
"is_hr": "true",
|
||
"user_img_path": "data:image/jpeg;base64,/9j/4AAQ...",
|
||
"output_format": "base64"
|
||
}'
|
||
```
|
||
|
||
```python
|
||
# Python 示例
|
||
import base64, requests
|
||
with open("user.jpg", "rb") as f:
|
||
img_b64 = base64.b64encode(f.read()).decode()
|
||
r = requests.post("http://127.0.0.1:8801/api/swapHair/v1", json={
|
||
"hair_id": "chang_tuoyuan",
|
||
"task_id": "req_001",
|
||
"is_hr": "true",
|
||
"user_img_path": "data:image/jpeg;base64," + img_b64,
|
||
"output_format": "base64",
|
||
}, timeout=300)
|
||
result = r.json() # {"state":0, "data":"<base64>", "msg":"success"}
|
||
with open("out.jpg", "wb") as f:
|
||
f.write(base64.b64decode(result["data"]))
|
||
```
|
||
|
||
### 响应
|
||
|
||
**成功（HTTP 200）：**
|
||
```json
|
||
{
|
||
"state": 0,
|
||
"msg": "success",
|
||
"data": "<base64 编码的结果图>" // output_format=url 时为 OSS 图片链接
|
||
}
|
||
```
|
||
|
||
**失败（HTTP 400）：**
|
||
```json
|
||
{
|
||
"state": -1,
|
||
"msg": "推理发型失败", // 具体错误信息，如：参数错误/用户图像加载失败/发型模板加载失败/推理发型失败
|
||
"data": "",
|
||
"task_id": "req_001"
|
||
}
|
||
```
|
||
|
||
---
|
||
|
||
## 四、换发型内部流程（9 个功能步骤）
|
||
|
||
理解内部流程有助于排查问题和性能优化。`hair_service_sd` 收到请求后依次执行：
|
||
|
||
| 步骤 | 功能 | 说明 | 耗时占比 |
|
||
|------|------|------|---------|
|
||
| 1 | 获取请求参数 | 解析参数，base64/URL 图片解码落盘 | <1% |
|
||
| 2 | 加载用户图像 | 复制到工作目录 | <1% |
|
||
| 3 | 加载发型模板图像 | 从 `upload_train_imgs/<hair_id>/` 读模板原图 | <1% |
|
||
| 4 | 检查遮挡眼睛 | 从 `hair_template_material/<hair_id>/` 读 `*_matting.png` 和 `.pkl`（关键点） | <1% |
|
||
| 5 | 检查发型材质 | 确认 `ref_hairstyle/<hair_id>/` 存在 | <1% |
|
||
| 6 | **推理发型** | 用 GAN 模型把用户图和 `ref_rgb_8uc3_768.png` 融合，生成初步换发图 | ~30% |
|
||
| 7 | 判断遮挡眼睛 | 检查头发是否遮挡眼睛（决定后续处理路径） | <1% |
|
||
| 8 | 处理发型区域 | 头发区域裁剪、mask 生成、透视变换到 576×768/1152×1536 | ~5% |
|
||
| 9 | **webui img2img** | 经 photo_service 调 webui，加载 LoRA 做 SD 重绘（核心画质步骤） | ~60% |
|
||
|
||
**关键：LoRA 加载机制**（步骤 9）
|
||
1. `hair_service_sd` 把裁剪好的头发区域图 + mask 发给 `photo_service`
|
||
2. `photo_service` 把 `train_material/<hair_id>/model/hairstyle_hd_lora.safetensors` 复制到 webui 的 `models/Lora/<hair_id>_hd.safetensors`（已存在则跳过）
|
||
3. `photo_service` 在 prompt 前注入 `<lora:<hair_id>_hd:1.0>`
|
||
4. webui 收到 img2img 请求，自动加载该 LoRA 做局部重绘
|
||
5. 结果返回 `hair_service_sd`，贴回原图对应位置
|
||
|
||
---
|
||
|
||
## 五、必需资源清单
|
||
|
||
换发型功能依赖以下资源。集成时需保证这些文件到位。
|
||
|
||
### 5.1 核心模型权重（全局共享，约 10G）
|
||
|
||
位置：`project/hair_service_sd/weights/`
|
||
|
||
| 文件/目录 | 用途 |
|
||
|-----------|------|
|
||
| `yolov5l.pt` | RetinaFace/人脸检测 |
|
||
| `Resnet50_Final.pth` | 人脸检测辅助 |
|
||
| `79999_iter.pth` | 关键点检测（MomocvFaceAlignment1K） |
|
||
| `master_hair_v8_onlyhair_nowarp_all_0709.pt` | 头发迁移 GAN（发型融合） |
|
||
| `zxm_v7_gen_bald_add_changed_hair_pairdata_5seg_0611.pt` | 秃头分割 + 发型生成 GAN |
|
||
| `parser_model.pth` | 头发抠图（matting） |
|
||
| `gender_models/` | 性别识别 |
|
||
| `keypoints/` | 关键点模型库 |
|
||
|
||
### 5.2 Stable Diffusion 推理资源（约 5G）
|
||
|
||
位置：`project/onediff/stable-diffusion-webui/models/`
|
||
|
||
| 文件/目录 | 用途 |
|
||
|-----------|------|
|
||
| `Stable-diffusion/v1-5-pruned-emaonly.safetensors` | **SD 1.5 底模**（必需，4G） |
|
||
| `Lora/<hair_id>_hd.safetensors` | 各发型的 LoRA（推理时自动生成，见 5.3） |
|
||
| `ESRGAN/4x-UltraSharp.pth` | 超分辨率（高清模式用） |
|
||
| `GFPGAN/` | 人脸修复（可选） |
|
||
|
||
### 5.3 单个发型的资源（每个发型约 160M）
|
||
|
||
**一个可用的发型需要 4 个目录的资源配套：**
|
||
|
||
```
|
||
# 1. 发型材质（换发型必需，由训练 step3 生成）
|
||
project/data/ref_hairstyle/<hair_id>/
|
||
├── config.json # {"gender":"girl","version":"20221206","ratio":"1"}
|
||
├── ref_rgb_8uc3_768.png # 发型 RGB 图（GAN 融合参考图，必需）
|
||
├── ref_matting_8uc3_768.png # 发型 mask
|
||
├── ref_matting_fg_8uc3_768.png # 发型前景
|
||
├── ref_baldseg_8uc3_768.png # 秃头分割
|
||
├── ref_landmark_f1k2_768.txt # 1000 个关键点
|
||
└── input_another_pose_hair_image.npy
|
||
|
||
# 2. LoRA 权重 + 训练图（换发型必需，由训练 step1+step2 生成）
|
||
data/train_material/<hair_id>/
|
||
├── model/
|
||
│ └── hairstyle_hd_lora.safetensors # LoRA 权重（~145M，SD 重绘用）
|
||
└── images/
|
||
└── 1_hairstyle/ # 训练图 + caption
|
||
├── 0001.png
|
||
├── 0001.txt # caption tag（webui 推理 prompt 用）
|
||
└── ... (约20张)
|
||
|
||
# 3. 模板材质（换发型必需，由训练 step4 生成）
|
||
project/data/hair_template_material/<hair_id>/
|
||
├── first##<uuid>.png # 发型模板原图
|
||
├── first##<uuid>_matting.png # 发型 matting（遮挡检测用）
|
||
└── first##<uuid>.pkl # 关键点（遮挡检测用）
|
||
|
||
# 4. 发型模板原图（换发型必需，由训练准备）
|
||
project/data/upload_train_imgs/<hair_id>/
|
||
└── first##<hair_id>.jpg # 发型模板原图（步骤3读取）
|
||
```
|
||
|
||
> **注意**：`images/1_hairstyle/` 目录名必须以数字开头（`1_hairstyle`），代码用 `os.listdir(images)[0]` 取子目录并校验数字开头。
|
||
|
||
### 5.4 路径配置
|
||
|
||
所有路径在 `project/hair_service_sd/config/configure.ini` 集中配置，迁移时改这里的路径即可：
|
||
|
||
```ini
|
||
[default]
|
||
modelDir = weights # 核心模型目录（相对 hair_service_sd）
|
||
hairstyleDir = /home/xsl/change_hair/project/data/ref_hairstyle # 发型材质（5.3-1）
|
||
haircolorDir= /home/xsl/change_hair/project/data/ref_haircolor # 换发色用
|
||
userDir=/home/xsl/change_hair/project/data/userImage # 用户图工作目录
|
||
tmp_dir=/home/xsl/change_hair/project/data/tmp # 临时目录
|
||
res_dir=/home/xsl/change_hair/project/data/res_dir # 结果输出目录
|
||
userInfo_dir=/home/xsl/change_hair/project/data/user_info # 用户关键点缓存
|
||
train_dir = /home/xsl/change_hair/data/train_material # LoRA+训练图（5.3-2）
|
||
refImgDir=/home/xsl/change_hair/project/data/refImage # 参考图缓存
|
||
hair_template_material_dir=/home/xsl/change_hair/project/data/hair_template_material # 5.3-3
|
||
upload_train_dir=/home/xsl/change_hair/project/data/upload_train_imgs # 5.3-4
|
||
version=online # online 用 0.0.0.0 地址
|
||
```
|
||
|
||
`photo_service` 的 webui/LoRA 路径在 `project/photo_service/lora_train_service_1.py` 第 41 行：
|
||
```python
|
||
webui_lora_dir = '/home/xsl/change_hair/project/onediff/stable-diffusion-webui/models/Lora'
|
||
```
|
||
|
||
---
|
||
|
||
## 六、集成步骤（接入新项目）
|
||
|
||
### 1. 部署服务端
|
||
按第二节启动 webui + photo_service + hair_service_sd 三个服务，确认 `http://<host>:8801` 可达。
|
||
|
||
### 2. 调用换发型接口
|
||
新项目只需调 `hair_service_sd` 的 `/api/swapHair/v1`，无需关心内部服务调用。
|
||
|
||
### 3. 准备发型资源
|
||
每个要支持的发型，必须备齐第五节 5.3 的 4 个目录资源。两种获取方式：
|
||
- **训练新发型**（见第七节）：用发型图训练，自动生成全部资源
|
||
- **复制现有发型**：从已部署环境复制 `<hair_id>` 的 4 个目录
|
||
|
||
### 4. 查询可用发型列表
|
||
调 `hair_grow_service` 的接口（需启动 8899）：
|
||
```
|
||
GET http://<host>:8899/api/hairstyles
|
||
→ {"state":0, "count":5, "data":[{"hair_id":"chang_tuoyuan","gender":"girl"}, ...]}
|
||
```
|
||
> 注意：该接口在 `hair_grow_swap.py:list_hairstyles()` 内有白名单过滤（当前只返回 5 个 chang_* 发型），集成时按需调整或去掉。
|
||
|
||
### 5. 获取发型预览图
|
||
```
|
||
GET http://<host>:8899/preview/<hair_id>
|
||
→ 直接返回预览图（优先 static/previews/<hair_id>.jpg，回退到 ref_rgb_8uc3_768.png）
|
||
```
|
||
|
||
---
|
||
|
||
## 七、训练新发型
|
||
|
||
### 7.1 一键训练脚本
|
||
|
||
```bash
|
||
cd /home/xsl/change_hair/project/hair_service_sd
|
||
python /home/xsl/change_hair/train_hairstyle_full.py \
|
||
--hair-id <新发型ID> \
|
||
--input <训练图目录> \ # 目录里放该发型的 jpg/png（1张即可，多张效果更好）
|
||
--gender <girl|boy> \
|
||
--template-img <模板图> # 用于生成材质和预览，通常是该发型最代表性的图
|
||
```
|
||
|
||
### 7.2 训练 5 个阶段
|
||
|
||
| 阶段 | 功能 | 产物 | 耗时 |
|
||
|------|------|------|------|
|
||
| step1 | 准备训练数据 | 抠图 + 数据增广（水平翻转 × 5 分辨率）→ `train_material/<hair_id>/images/1_hairstyle/` 约 20 张 png+txt | ~1 分钟 |
|
||
| step2 | LoRA 训练 | kohya 训练 1500 步 → `train_material/<hair_id>/model/hairstyle_hd_lora.safetensors` | ~10 分钟（单卡） |
|
||
| step3 | 生成 ref 材质 | 调 hair_service_sd 回调 → `ref_hairstyle/<hair_id>/` 6 个文件 | ~1 秒 |
|
||
| step4 | 生成模板材质 | → `hair_template_material/<hair_id>/` 3 个文件 | ~3 秒 |
|
||
| step5 | 生成预览图 | → `hair_grow_service/static/previews/<hair_id>.jpg` | ~10 秒 |
|
||
|
||
### 7.3 批量并行训练（多个发型）
|
||
|
||
用 `train_lora_parallel.py` 直接调 kohya，支持并发：
|
||
|
||
```python
|
||
# 编辑 train_lora_parallel.py 顶部的 HAIRSTYLES 列表
|
||
HAIRSTYLES = ["新发型1", "新发型2", "新发型3"]
|
||
PARALLEL = 2 # 并发数（按显存调整，5090 32G 建议 2）
|
||
|
||
# step1 数据需先用 train_hairstyles_parallel.py 或逐个跑 step1 准备好
|
||
python train_lora_parallel.py
|
||
```
|
||
|
||
### 7.4 训练注意事项
|
||
- **hair_id 命名**：只能用英文/数字/下划线，不能含中文或特殊字符（会作目录名、文件名、LoRA 名）
|
||
- **单图训练**：1 张图经增广出 ~20 张训练图，LoRA 能用但泛化有限；多图（不同角度/光照）效果更好
|
||
- **GPU 显存**：训练峰值约 13G/卡，推理约 6G。并行训练需控制并发数避免 OOM
|
||
- **训练前提**：webui + photo_service + hair_service_sd 必须在运行（step3 回调依赖）
|
||
|
||
---
|
||
|
||
## 八、环境依赖
|
||
|
||
### Conda 环境
|
||
|
||
| 环境 | Python | 用途 | 关键依赖 |
|
||
|------|--------|------|---------|
|
||
| `my_hair` | 3.10 | hair_service_sd / hair_grow_service / 训练 step1 | torch, cv2, flask, gevent, oss2 |
|
||
| `sdwebui` | 3.10 | webui | torch, gradio, transformers（AUTOMATIC1111 全套） |
|
||
| `py310` | 3.10 | photo_service | torch, flask, requests |
|
||
| `kohya` | 3.10 | LoRA 训练（kohya_ss） | torch, accelerate, safetensors |
|
||
|
||
### 关键兼容性处理（迁移时注意）
|
||
|
||
1. **PyTorch 2.6+ 的 `weights_only`**：加载旧权重（yolov5l.pt 等）需 `torch.load(..., weights_only=False)`。代码已在 `run_copy_cost_colorb64.py` 和 `hair_grow_service/app.py` 顶部做了全局 patch。
|
||
|
||
2. **OSS 上传（可选）**：`output_format=url` 时需配 OSS 凭证环境变量：
|
||
```bash
|
||
export OSS_TEST_ACCESS_KEY_ID=<...>
|
||
export OSS_TEST_ACCESS_KEY_SECRET=<...>
|
||
export OSS_TEST_BUCKET=<...>
|
||
export OSS_TEST_ENDPOINT=<...>
|
||
```
|
||
用 `output_format=base64` 则完全不需要 OSS。
|
||
|
||
3. **webui xformers**：kohya 训练命令用 `--sdpa`（PyTorch 原生），不要用 `--xformers`（若 xformers 与 torch 版本不匹配会报 CUDA 错误）。
|
||
|
||
---
|
||
|
||
## 九、故障排查
|
||
|
||
| 现象 | 可能原因 | 排查 |
|
||
|------|---------|------|
|
||
| 换发型返回 `推理发型失败` | webui 没启动 / LoRA 文件缺失 | 检查 webui 57860 端口；检查 `train_material/<hair_id>/model/` 是否有 LoRA |
|
||
| 换发型返回 `发型模板加载失败` | `upload_train_imgs/<hair_id>/` 缺 `first##` 开头的文件 | 检查 5.3-4 目录 |
|
||
| 换发型返回 `检查遮挡失败` | `hair_template_material/<hair_id>/` 缺 matting/pkl | 检查 5.3-3 目录 |
|
||
| 预览图 404 | `ref_hairstyle/<hair_id>/ref_rgb_8uc3_768.png` 不存在 | 检查 5.3-1 目录 |
|
||
| 发型列表少发型 | `hair_grow_swap.py:list_hairstyles()` 白名单过滤 | 调整 `_VISIBLE_HAIRSTYLES` 集合 |
|
||
| webui 启动报 xformers 错 | xformers 与 torch 版本不匹配 | webui 启动去掉 `--xformers`，或用 `--sdpa` |
|
||
| 训练报 CUDA error | 显存不足或 xformers 问题 | 训练命令用 `--sdpa`；降低并发数 |
|
||
|
||
### 日志位置
|
||
- webui：`project/logs/webui.log`
|
||
- photo_service：`project/logs/photo_service.log`
|
||
- hair_service_sd：`project/logs/hair_service.log`
|
||
- hair_grow_service：`project/logs/hair_grow_service.log`
|
||
- 训练：`project/logs/train_<hair_id>.log`
|
||
|
||
---
|
||
|
||
## 十、目录结构总览
|
||
|
||
```
|
||
/home/xsl/change_hair/
|
||
├── project/
|
||
│ ├── hair_service_sd/ # 换发型主服务
|
||
│ │ ├── weights/ # 核心模型权重（10G）
|
||
│ │ ├── config/configure.ini # 路径配置
|
||
│ │ ├── run_copy_cost_colorb64.py # 主入口（API 定义）
|
||
│ │ ├── gen_super_image.py # webui img2img 调用
|
||
│ │ └── core/ # 发型推理核心逻辑
|
||
│ ├── photo_service/ # LoRA 调度服务
|
||
│ │ └── lora_train_service_1.py # 训练调度 + 推理 LoRA 切换
|
||
│ ├── onediff/ # SD webui（推理引擎，14G）
|
||
│ │ └── stable-diffusion-webui/models/
|
||
│ │ ├── Stable-diffusion/ # SD 底模
|
||
│ │ └── Lora/ # 各发型 LoRA（自动生成）
|
||
│ ├── kohya_ss_home/ # LoRA 训练框架（13G）
|
||
│ ├── data/ # 数据目录（由 configure.ini 指向）
|
||
│ │ ├── ref_hairstyle/ # 发型材质（5.3-1）
|
||
│ │ ├── hair_template_material/ # 模板材质（5.3-3）
|
||
│ │ ├── upload_train_imgs/ # 发型模板原图（5.3-4）
|
||
│ │ └── ...
|
||
│ └── logs/
|
||
├── data/
|
||
│ └── train_material/ # LoRA + 训练图（5.3-2）
|
||
├── hair_grow_service/ # 前端页面（可选）
|
||
├── train_hairstyle_full.py # 单发型一键训练
|
||
├── train_lora_parallel.py # 批量并行训练
|
||
├── start_all.sh # 服务启停脚本
|
||
└── docs/换发型集成文档.md # 本文档
|
||
```
|