# 换发型功能集成文档

本文档说明"换发型"功能的完整工作流程、服务依赖、资源结构，供其他项目集成参考。

---

## 一、功能概述

**换发型**：用户上传一张人像照片，选择一个目标发型（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            # 本文档
```