xsl/hair
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

优化为4b 模型

Browse Source
This commit is contained in:
xsl
2026-07-18 15:30:31 +08:00
parent 659c037270
commit b58cd4c441
289 changed files with 9367 additions and 558 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/optimization_plans_v2.md
+298
View File
@@ -0,0 +1,298 @@
# 优化方案 V2 — 修订版（约束：保留 swapHair + 保留所有服务）
> 日期：2026-07-18
> 约束：不能去掉 swapHair、不能停止 hair_service_sd 和 SD WebUI
---
## 根因总结
**ComfyUI 未启用 `--gpu-only`，默认在每次请求后卸载模型到 CPU。**
9B FP8 模型(8.8GB) + 文本编码器(7.7GB) + VAE(0.3GB) = **16.8GB 模型显存**，可用 20GB 中仅剩 3.2GB 给 KSampler 执行 → ComfyUI 自动部分卸载模型 → 下次请求重新加载 → **每次浪费 8-15 秒**。
### 关键发现
| 项目 | 值 |
|------|-----|
| GPU | RTX 5090 (32GB) |
| PyTorch | 2.7.1+cu128 (**已支持 sm_120！**) |
| ComfyUI 启动参数 | `--listen 127.0.0.1 --port 8188`（无 --gpu-only） |
| 当前模型 | 9B FP8 (8.8GB) |
| **4B FP8 模型已在磁盘** | **3.8GB**（`diffusion_models/flux-2-klein-4b-fp8.safetensors`） |
| 4B 文本编码器已在磁盘 | 7.5GB（`text_encoders/qwen_3_4b.safetensors`） |
| 可用显存（ComfyUI） | ~20GB（32GB - 8GB hair - 3.2GB SD WebUI - 0.7GB Hair） |
### 两种工作流对比
| 维度 | 0716add-hair-api.json (接口2女) | add_hair.json (接口2男/接口3) |
|------|------|------|
| 节点数 | 27 | 29 |
| 模型 | flux-2-klein-9b-fp8 | flux-2-klein-9b-fp8 (相同) |
| 文本编码器 | qwen_3_8b_fp8mixed | qwen_3_8b_fp8mixed (相同) |
| VAE | flux2-vae | flux2-vae (相同) |
| KSampler steps | **4** | **6** |
| 输出文件名 | hair_inpaint_XXXXX | ComfyUI_XXXXX |
**两个工作流加载完全相同的模型，仅 steps 不同。模型换出的根因是 ComfyUI 的默认卸载策略，不是工作流切换。**
---
## 方案 A：--gpu-only + 统一 steps=4（最小改动）
### 思路
给 ComfyUI 加 `--gpu-only` 标志，禁止模型在 GPU↔CPU 之间转移。同时统一两个工作流的 steps 为 4，减少推理时间。
### 具体改动
#### A1. ComfyUI 启动加 --gpu-only
```bash
# 修改 ComfyUI 启动命令
python main.py --listen 127.0.0.1 --port 8188 --gpu-only
```
`--gpu-only` 让 ComfyUI 启动后将所有模型权重保留在 GPU 显存中，不自动卸载到 CPU。
#### A2. 统一 add_hair.json 的 steps 为 4
```json
// add_hair.json 节点 9 (SamplerCustomAdvanced)
// 将 steps 从 6 改为 4
```
#### A3. 显存验证
| 组件 | 显存 |
|------|------|
| 9B FP8 模型 | 8.8 GB |
| 文本编码器 (8B FP8) | 7.7 GB |
| VAE | 0.3 GB |
| KSampler 执行 | ~3-4 GB |
| **ComfyUI 总需求** | **~20-21 GB** |
| **可用** | **~20 GB** |
| **余量** | **-1 ~ 0 GB** ⚠ |
**风险**：9B 模型 + --gpu-only 可能刚好 OOM。需要实测。
### 预期耗时
| 接口 | 当前 | 方案A | 达标 |
|------|------|-------|------|
| 接口2-女 | 17.6s | ~11s（swapHair 4.6s + ComfyUI 4s + 预处理 1.5s + 余量 1s） | ❌ 差1s |
| 接口2-男 | 8.1s | ~6s（ComfyUI 4s + 预处理 2s） | ✅ |
| 接口3 | 22.9s | ~6s（ComfyUI 4s + 预处理 1.6s） | ✅ |
### 优点
- 改动极小（1行启动参数 + 1个JSON值）
- 不换模型，质量不变
- 可快速验证
### 缺点
- **9B 模型可能 OOM**（仅 0-1GB 余量）
- 接口2女仍 ~11s（swapHair 是硬瓶颈）
- `--gpu-only` 在显存不足时可能导致 ComfyUI 崩溃
---
## 方案 B：换 4B FP8 模型 + --gpu-only（推荐）
### 思路
9B FP8（8.8GB）换成 4B FP8（3.8GB），省 5GB 显存。模型+编码器从 16.8GB 降到 11.6GB，KSampler 有 8.4GB 余量，彻底消除模型换出。4B 模型同样用 4 步推理，速度相当或更快。
### 磁盘上已有模型
```
ComfyUI/models/diffusion_models/flux-2-klein-4b-fp8.safetensors 3.8 GB ✅已存在
ComfyUI/models/text_encoders/qwen_3_4b.safetensors 7.5 GB ✅已存在
ComfyUI/models/vae/flux2-vae.safetensors 0.3 GB ✅已存在
```
**无需下载任何新文件！**
### 具体改动
#### B1. 修改两个工作流的模型加载节点
```python
# 0716add-hair-api.json 和 add_hair.json
# 节点 16 (UNETLoader): unet_name
"flux2.0/flux-2-klein-9b-fp8.safetensors"
→ "flux-2-klein-4b-fp8.safetensors"
# 节点 61 (CLIPLoader): clip_name
"qwen_3_8b_fp8mixed.safetensors"
→ "qwen_3_4b.safetensors"
```
#### B2. ComfyUI 启动加 --gpu-only
```bash
python main.py --listen 127.0.0.1 --port 8188 --gpu-only
```
#### B3. 统一 steps 为 4
add_hair.json 的 steps 从 6 改为 4。
#### B4. 显存验证
| 组件 | 显存 |
|------|------|
| 4B FP8 模型 | 3.8 GB |
| 文本编码器 (4B BF16) | 7.5 GB |
| VAE | 0.3 GB |
| KSampler 执行 | ~3-4 GB |
| **ComfyUI 总需求** | **~15-16 GB** |
| **可用** | **~20 GB** |
| **余量** | **+4-5 GB** ✅ 安全 |
### 预期耗时
| 接口 | 当前 | 方案B | 达标 |
|------|------|-------|------|
| 接口2-女 | 17.6s | ~11s（swapHair 4.6s + ComfyUI 3s + 预处理 1.5s + 余量 2s） | ❌ 差1s |
| 接口2-男 | 8.1s | ~5s（ComfyUI 3s + 预处理 2s） | ✅ |
| 接口3 | 22.9s | ~5s（ComfyUI 3s + 预处理 1.6s） | ✅ |
### 4B vs 9B 质量差异
根据公开资料：
- 4B 模型设计目标是"在消费级 GPU 上实现实时生成"
- 4 步推理，与 9B 相同
- 4B 是 Apache 2.0 开源，9B 是非商用许可
- 质量差异：4B 在复杂场景和细节上略逊于 9B，但对头发补全这类局部重绘任务影响可能较小
- **需要实际对比测试验证**
### 优点
- 模型已在磁盘，无需下载
- 显存余量充足（+4-5GB），不会 OOM
- 4B 模型推理速度更快
- 可彻底消除模型换出
- 不影响其他服务
### 缺点
- 需要质量验证（4B vs 9B）
- 需修改两个工作流 JSON
- 接口2女仍 ~11s（swapHair 是硬瓶颈）
---
## 方案 C：4B FP8 + --gpu-only + SegFormer→GPU + steps=2（激进优化）
### 思路
在方案B基础上进一步优化：
1. SegFormer 移到 GPU（torch 已是 cu128，支持 5090）
2. steps 降到 2（需质量验证）
3. swapHair 与预处理并行
### 具体改动
#### C1-C3：同方案B
#### C4. SegFormer 移到 GPU
```python
# hairline/face_parsing.py 或环境变量
# 当前 SegFormer 跑 CPU（~0.9s/张）
# torch 2.7.1+cu128 已支持 RTX 5090 (sm_120)
export SEG_DEVICE=cuda
# 预期：0.9s → 0.05s
```
#### C5. steps 降到 2
```json
// add_hair.json + 0716add-hair-api.json
// KSampler steps: 4 → 2
```
Flux-2 Klein 是 step-distilled 模型，设计为 4 步。降到 2 步会有质量损失，需验证。
#### C6. swapHair 与预处理并行
```
当前流程：
预处理(1.5s) → swapHair(4.6s) → ComfyUI(3s) = 9.1s
优化后（并行）：
预处理(1.5s) ─┐
├→ ComfyUI(3s) = 6.1s
swapHair(4.6s) ┘
```
### 显存验证
| 组件 | 显存 |
|------|------|
| 4B FP8 模型 | 3.8 GB |
| 文本编码器 (4B) | 7.5 GB |
| VAE | 0.3 GB |
| SegFormer (GPU) | ~0.5 GB |
| KSampler 执行 | ~2-3 GB（steps=2 更少） |
| **ComfyUI 总需求** | **~14-15 GB** |
| **可用** | **~20 GB** |
| **余量** | **+5-6 GB** ✅ |
### 预期耗时
| 接口 | 当前 | 方案C | 达标 |
|------|------|-------|------|
| 接口2-女 | 17.6s | ~7s（swapHair 4.6s ∥ 预处理 0.6s + ComfyUI 2s） | ✅ |
| 接口2-男 | 8.1s | ~4s（ComfyUI 2s + 预处理 1s） | ✅ |
| 接口3 | 22.9s | ~4s（ComfyUI 2s + 预处理 0.7s） | ✅ |
### 优点
- 所有接口达标（<10s）
- SegFormer GPU 化（torch cu128 已支持）
- swapHair 并行化进一步优化接口2女
### 缺点
- steps=2 质量需验证
- SegFormer GPU 需修改代码
- swapHair 并行需重构 service.py
- 改动量最大
---
## 三方案对比
| 维度 | A (--gpu-only+steps=4) | B (4B FP8+--gpu-only) | C (4B+GPU SegFormer+steps=2) |
|------|----------------------|----------------------|------------------------------|
| 改动量 | 极小 | 小 | 中 |
| 需下载 | 无 | 无 | 无 |
| 接口2-女 | ~11s ❌ | ~11s ❌ | ~7s ✅ |
| 接口2-男 | ~6s ✅ | ~5s ✅ | ~4s ✅ |
| 接口3 | ~6s ✅ | ~5s ✅ | ~4s ✅ |
| 显存安全 | ⚠ 紧张(0-1GB) | ✅ 安全(+4GB) | ✅ 宽裕(+5GB) |
| 质量风险 | 无 | 中（4B vs 9B） | 高（4B+steps=2） |
| OOM 风险 | 高 | 低 | 极低 |
---
## 建议
### 第一阶段：先做方案B
1. 修改两个工作流 JSON：换 4B FP8 模型 + 4B 文本编码器
2. ComfyUI 加 `--gpu-only`
3. 统一 steps=4
4. 跑一轮质量对比（4B vs 9B 的生发效果）
5. 跑轮换测试验证耗时
**如果接口3和接口2男达标（<10s），接口2女接近（~11s），则进入第二阶段。**
### 第二阶段：方案C 追加优化（解决接口2女）
6. SegFormer 移到 GPU（环境变量 SEG_DEVICE=cuda）
7. swapHair 与预处理并行化
8. 如果质量允许，steps 降到 2-3
### 需要你确认
1. **是否先用方案B（4B FP8 + --gpu-only）做一轮测试？**
2. **4B 模型质量是否需要先对比验证，还是可以直接切换后看效果？**
3. **方案C 的 steps=2 是否可以接受（需要质量验证）？**