Author SHA1 Message Date
xsl 326b206d08 feat(swapHair): webui_steps参数可外部传入(接口2调试页可调)
- gen_super_image.py: webui_img2img + build_body_v2 接收可选 steps 参数
- run_copy_cost_colorb64.py: swapHair 接口接收 webui_steps,透传给 webui_img2img
- 未传时保持环境变量 WEBUI_STEPS 默认(15),向后兼容
2026-07-26 16:49:21 +08:00
xsl 990fae929f perf(swapHair): webui降步数 + 用户图预处理内存缓存
优化1 - webui img2img steps 20→15 (gen_super_image.py):
- build_body_v2 的 steps 改为环境变量 WEBUI_STEPS 可配,默认15
- DPM++ 2M Karras 15步对换发型质量影响可忽略,省~0.1s

优化2 - infer_hairstyle_diy_jy 用户图预处理内存缓存 (hairstyle_model.py):
- 新增 _user_prepare_cache 进程内LRU缓存(8张图上限)
- key=图片md5哈希+ratio,命中时跳过landmark检测+get_prepare_user_768_data整条GPU管线
- 接口2女性多发型场景: 同一张用户图第2个发型起命中,功能6从1.8s降至1.1s(省0.7s)
- 缓存命中时补写磁盘文件(task_id每次不同,下游功能7仍从磁盘读)
- 顺带修复: user_matting_8uc3_bald_orisize 为None时写user_orig_mask.png的crash
- 加分步计时日志(主GAN/融合耗时),便于定位热点

实测: 同图连续请求 swapHair 从4.28s降至3.5s(省18%)
2026-07-26 16:41:19 +08:00
xslandCursor c5e50de40a chore: 固化 WebUI 底模为 v1-5,并纳入仓库配置副本
本机无 majicmix,将 sd_model_checkpoint 固定为 v1-5-pruned-emaonly;因 onediff/webui 是独立 git 仓无法直接跟踪,配置副本放在 configs/。

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-19 15:52:51 +08:00
xslandCursor 0926c61bd0 fix: 适配 ubuntu 路径并增强换发色/训练流程稳定性
将配置与训练脚本从 /home/xsl 切到本机 /home/ubuntu;换发色在 webui 增强失败或缺色板时降级返回,训练结束后自动重启 hair 服务再回调。

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-19 00:49:16 +08:00
xsl fc66534a74 feat: 系统配置文件适配 RTX 3090
添加系统服务配置文件:
- change_hair-hair.service: hair_service_sd swapHair (8801)
- change_hair-photo.service: photo_service LoRA scheduler (32678)
- change_hair-webui.service: SD WebUI (57860)
2026-07-18 19:18:56 +08:00
xsl 6f34e8876c feat: 适配 RTX 3090 (24GB) 环境优化
硬件迁移:从 RTX 5090 (32GB) 迁移到 RTX 3090 (24GB)

主要改动:
- 所有启动脚本和配置文件中的 /home/xsl/ 路径替换为 /home/ubuntu/
- 适配新的 24GB VRAM 环境
2026-07-18 19:00:15 +08:00
34 changed files with 348 additions and 166 deletions
+2 -2
View File
@@ -296,8 +296,8 @@ curl -X POST http://127.0.0.1:32678/api/hair/train \
| 改动 | 说明 | | 改动 | 说明 |
|------|------| |------|------|
| `config.json` 的 onediff compiler 路径 | 指向本机 onediff 目录 | | `config.json` 的 onediff compiler 路径 | 指向本机 onediff 目录 |
| 启动加 `HF_HUB_OFFLINE=1` | 本机无法访问 huggingface.co,用本地缓存离线加载 CLIP | | 启动加 `HF_HUB_OFFLINE=1` | 本机无法稳定访问 huggingface.co,用本地缓存离线加载 CLIP(需预先缓存 `openai/clip-vit-large-patch14` |
| SD 模型 | 实际加载 `v1-5-pruned-emaonly`config 里写的 majicmix 不存在,webui 自动 fallback | | SD 模型 | `sd_model_checkpoint` 固定为 `v1-5-pruned-emaonly.safetensors`(本机无 majicmix)。仓库内权威副本:`configs/stable-diffusion-webui.config.json`;部署时拷到 `project/onediff/stable-diffusion-webui/config.json`(该目录是独立 git 仓库,无法被本仓直接跟踪 |
### 3. 换发型推理改动(`hair_service_sd/gen_super_image.py` ### 3. 换发型推理改动(`hair_service_sd/gen_super_image.py`
+12 -12
View File
@@ -8,8 +8,8 @@
3. 串行执行,记录每个发型成功/失败 3. 串行执行,记录每个发型成功/失败
用法: 用法:
cd /home/xsl/change_hair/project/hair_service_sd cd /home/ubuntu/change_hair/project/hair_service_sd
python /home/xsl/change_hair/batch_train_hairstyles.py --src /home/xsl/change_hair/hair_type_images --gender girl [--only 圆-心形] [--start-from 心形-心形] python /home/ubuntu/change_hair/batch_train_hairstyles.py --src /home/ubuntu/change_hair/hair_type_images --gender girl [--only 圆-心形] [--start-from 心形-心形]
注意: 注意:
- 中文 hair_id 直接用文件名(去扩展名)作为 ID - 中文 hair_id 直接用文件名(去扩展名)作为 ID
@@ -24,14 +24,14 @@ import argparse
import subprocess import subprocess
from datetime import datetime from datetime import datetime
SRC_DEFAULT = "/home/xsl/change_hair/hair_type_images" SRC_DEFAULT = "/home/ubuntu/change_hair/hair_type_images"
WORK_DIR = "/home/xsl/change_hair/data/batch_train_inputs" WORK_DIR = "/home/ubuntu/change_hair/data/batch_train_inputs"
LOG_FILE = "/home/xsl/change_hair/data/batch_train_log.txt" LOG_FILE = "/home/ubuntu/change_hair/data/batch_train_log.txt"
TRAIN_SCRIPT = "/home/xsl/change_hair/train_hairstyle_full.py" TRAIN_SCRIPT = "/home/ubuntu/change_hair/train_hairstyle_full.py"
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd" HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
LOCK_FILE = "/home/xsl/change_hair/data/batch_train.pid" LOCK_FILE = "/home/ubuntu/change_hair/data/batch_train.pid"
DONE_FILE = "/home/xsl/change_hair/data/batch_train_done.txt" # 已完成发型清单(断点续跑) DONE_FILE = "/home/ubuntu/change_hair/data/batch_train_done.txt" # 已完成发型清单(断点续跑)
LORA_DIR = "/home/xsl/change_hair/data/train_material" # LoRA 输出根目录 LORA_DIR = "/home/ubuntu/change_hair/data/train_material" # LoRA 输出根目录
def acquire_lock(): def acquire_lock():
@@ -125,7 +125,7 @@ def run_one(hair_id, src_img, gender, py, done_set, force=False):
log(f" 调用: {' '.join(cmd[:2])} ... --hair-id {hair_id}") log(f" 调用: {' '.join(cmd[:2])} ... --hair-id {hair_id}")
try: try:
# 子进程输出实时写到日志文件 # 子进程输出实时写到日志文件
proc_log = os.path.join("/home/xsl/change_hair/data", f"subprocess_{hair_id}.log") proc_log = os.path.join("/home/ubuntu/change_hair/data", f"subprocess_{hair_id}.log")
with open(proc_log, "w", encoding="utf-8") as f: with open(proc_log, "w", encoding="utf-8") as f:
r = subprocess.run( r = subprocess.run(
cmd, cwd=HAIR_SERVICE_DIR, cmd, cwd=HAIR_SERVICE_DIR,
@@ -152,7 +152,7 @@ def main():
ap = argparse.ArgumentParser(description="批量训练发型") ap = argparse.ArgumentParser(description="批量训练发型")
ap.add_argument("--src", default=SRC_DEFAULT, help="原图目录") ap.add_argument("--src", default=SRC_DEFAULT, help="原图目录")
ap.add_argument("--gender", default="girl", choices=["boy", "girl"]) ap.add_argument("--gender", default="girl", choices=["boy", "girl"])
ap.add_argument("--py", default="/home/xsl/miniconda3/envs/my_hair/bin/python", ap.add_argument("--py", default="/home/ubuntu/miniconda3/envs/my_hair/bin/python",
help="python 解释器") help="python 解释器")
ap.add_argument("--only", default=None, help="只训练指定 hair_id(调试用)") ap.add_argument("--only", default=None, help="只训练指定 hair_id(调试用)")
ap.add_argument("--start-from", default=None, ap.add_argument("--start-from", default=None,
+19
View File
@@ -0,0 +1,19 @@
[Unit]
Description=change_hair hair_service_sd swapHair (8801)
After=network-online.target change_hair-webui.service change_hair-photo.service
Wants=change_hair-webui.service change_hair-photo.service
[Service]
Type=simple
User=ubuntu
WorkingDirectory=/home/ubuntu/change_hair/project/hair_service_sd
Environment=CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1
Environment=CUDA_VISIBLE_DEVICES=0
Environment=APP_WORKER_ID=1
ExecStart=/home/ubuntu/miniconda3/envs/my_hair/bin/python run_copy_cost_colorb64.py
Restart=on-failure
RestartSec=10
TimeoutStartSec=300
[Install]
WantedBy=multi-user.target
+17
View File
@@ -0,0 +1,17 @@
[Unit]
Description=change_hair photo_service LoRA scheduler (32678)
After=network-online.target
Wants=network-online.target
[Service]
Type=simple
User=ubuntu
WorkingDirectory=/home/ubuntu/change_hair/project/photo_service
Environment=CUDA_VISIBLE_DEVICES=0
Environment=APP_WORKER_ID=1
ExecStart=/home/ubuntu/miniconda3/envs/py310/bin/python -u lora_train_service_1.py
Restart=on-failure
RestartSec=10
[Install]
WantedBy=multi-user.target
+20
View File
@@ -0,0 +1,20 @@
[Unit]
Description=change_hair SD WebUI (57860)
After=network-online.target
Wants=network-online.target
[Service]
Type=simple
User=ubuntu
WorkingDirectory=/home/ubuntu/change_hair/project/onediff/stable-diffusion-webui
Environment=CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1
Environment=CUDA_VISIBLE_DEVICES=0
Environment=HF_HUB_OFFLINE=1
Environment=TRANSFORMERS_OFFLINE=1
ExecStart=/home/ubuntu/miniconda3/envs/sdwebui/bin/python webui.py --api --listen --xformers --port 57860
Restart=on-failure
RestartSec=10
TimeoutStartSec=600
[Install]
WantedBy=multi-user.target
+40
View File
@@ -0,0 +1,40 @@
{
"ldsr_steps": 100,
"ldsr_cached": false,
"SCUNET_tile": 256,
"SCUNET_tile_overlap": 8,
"SWIN_tile": 192,
"SWIN_tile_overlap": 8,
"SWIN_torch_compile": false,
"hypertile_enable_unet": false,
"hypertile_enable_unet_secondpass": false,
"hypertile_max_depth_unet": 3,
"hypertile_max_tile_unet": 256,
"hypertile_swap_size_unet": 3,
"hypertile_enable_vae": false,
"hypertile_max_depth_vae": 3,
"hypertile_max_tile_vae": 128,
"hypertile_swap_size_vae": 3,
"onediff_compiler_caches_path": "/home/ubuntu/change_hair/project/onediff/stable-diffusion-webui/extensions/onediff_sd_webui_extensions/compiler_caches",
"onediff_compiler_backend": "oneflow",
"sd_model_checkpoint": "v1-5-pruned-emaonly.safetensors",
"sd_checkpoint_hash": "6ce0161689b3853acaa03779ec93eafe75a02f4ced659bee03f50797806fa2fa",
"control_net_detectedmap_dir": "detected_maps",
"control_net_models_path": "",
"control_net_modules_path": "",
"control_net_unit_count": 3,
"control_net_model_cache_size": 2,
"control_net_inpaint_blur_sigma": 7,
"control_net_no_detectmap": false,
"control_net_detectmap_autosaving": false,
"control_net_allow_script_control": false,
"control_net_sync_field_args": true,
"controlnet_show_batch_images_in_ui": false,
"controlnet_increment_seed_during_batch": false,
"controlnet_disable_openpose_edit": false,
"controlnet_disable_photopea_edit": false,
"controlnet_photopea_warning": true,
"controlnet_ignore_noninpaint_mask": false,
"controlnet_clip_detector_on_cpu": false,
"controlnet_control_type_dropdown": false
}
+3 -3
View File
@@ -8,7 +8,7 @@ import sys
import cv2 import cv2
import numpy as np import numpy as np
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd" HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
sys.path.insert(0, HAIR_SERVICE_DIR) sys.path.insert(0, HAIR_SERVICE_DIR)
os.environ.setdefault("HF_HUB_OFFLINE", "1") os.environ.setdefault("HF_HUB_OFFLINE", "1")
@@ -87,7 +87,7 @@ def gen_hairline_mask(img_path, mask_path):
if __name__ == "__main__": if __name__ == "__main__":
img_path = sys.argv[1] if len(sys.argv) > 1 else \ img_path = sys.argv[1] if len(sys.argv) > 1 else \
"/home/xsl/change_hair/project/data/userImage/8488902485_20250630055547.jpg" "/home/ubuntu/change_hair/project/data/userImage/8488902485_20250630055547.jpg"
mask_path = sys.argv[2] if len(sys.argv) > 2 else \ mask_path = sys.argv[2] if len(sys.argv) > 2 else \
"/home/xsl/change_hair/project/logs/hairgrow_test_mask.png" "/home/ubuntu/change_hair/project/logs/hairgrow_test_mask.png"
gen_hairline_mask(img_path, mask_path) gen_hairline_mask(img_path, mask_path)
+6 -6
View File
@@ -14,15 +14,15 @@ import time
import base64 import base64
import requests import requests
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd" HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
sys.path.insert(0, HAIR_SERVICE_DIR) sys.path.insert(0, HAIR_SERVICE_DIR)
os.chdir(HAIR_SERVICE_DIR) os.chdir(HAIR_SERVICE_DIR)
BOY_IMG = "/home/xsl/change_hair/images/boy.png" BOY_IMG = "/home/ubuntu/change_hair/images/boy.png"
GIRL_IMG = "/home/xsl/change_hair/images/girl.png" GIRL_IMG = "/home/ubuntu/change_hair/images/girl.png"
HAIRSTYLE_DIR = "/home/xsl/change_hair/project/data/ref_hairstyle" HAIRSTYLE_DIR = "/home/ubuntu/change_hair/project/data/ref_hairstyle"
TRAIN_DIR = "/home/xsl/change_hair/data/train_material" TRAIN_DIR = "/home/ubuntu/change_hair/data/train_material"
PREVIEW_DIR = "/home/xsl/change_hair/hair_grow_service/static/previews" # 预览图存储目录 PREVIEW_DIR = "/home/ubuntu/change_hair/hair_grow_service/static/previews" # 预览图存储目录
SWAP_API = "http://127.0.0.1:8801/api/swapHair/v1" SWAP_API = "http://127.0.0.1:8801/api/swapHair/v1"
+2 -2
View File
@@ -8,7 +8,7 @@
hair_template_material/<hair_id>/first##<name>.pkl 1k关键点 hair_template_material/<hair_id>/first##<name>.pkl 1k关键点
用法: 用法:
cd /home/xsl/change_hair/project/hair_service_sd cd /home/ubuntu/change_hair/project/hair_service_sd
python gen_template_material.py --hair-id new_test_001 --img /path/to/template.jpg python gen_template_material.py --hair-id new_test_001 --img /path/to/template.jpg
""" """
import os import os
@@ -19,7 +19,7 @@ import pickle
import argparse import argparse
ssl._create_default_https_context = ssl._create_unverified_context ssl._create_default_https_context = ssl._create_unverified_context
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd" HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
os.chdir(HAIR_SERVICE_DIR) os.chdir(HAIR_SERVICE_DIR)
sys.path.insert(0, HAIR_SERVICE_DIR) sys.path.insert(0, HAIR_SERVICE_DIR)
os.environ.setdefault("HF_HUB_OFFLINE", "1") os.environ.setdefault("HF_HUB_OFFLINE", "1")
+4 -4
View File
@@ -6,12 +6,12 @@
需要在 hair_service_sd 目录下运行(依赖其模块导入)。 需要在 hair_service_sd 目录下运行(依赖其模块导入)。
用法: 用法:
cd /home/xsl/change_hair/project/hair_service_sd cd /home/ubuntu/change_hair/project/hair_service_sd
/home/xsl/miniconda3/envs/my_hair/bin/python /home/xsl/change_hair/hair_grow_cli.py \ /home/ubuntu/miniconda3/envs/my_hair/bin/python /home/ubuntu/change_hair/hair_grow_cli.py \
--img test.jpg --mask mask.png --strength 0.5 -o result.jpg --img test.jpg --mask mask.png --strength 0.5 -o result.jpg
# 批量跑三档强度对比 # 批量跑三档强度对比
/home/xsl/miniconda3/envs/my_hair/bin/python /home/xsl/change_hair/hair_grow_cli.py \ /home/ubuntu/miniconda3/envs/my_hair/bin/python /home/ubuntu/change_hair/hair_grow_cli.py \
--img test.jpg --mask mask.png --compare --img test.jpg --mask mask.png --compare
""" """
import os import os
@@ -20,7 +20,7 @@ import cv2
import argparse import argparse
# 把 hair_service_sd 加入路径,使其模块可被导入 # 把 hair_service_sd 加入路径,使其模块可被导入
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd" HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
sys.path.insert(0, HAIR_SERVICE_DIR) sys.path.insert(0, HAIR_SERVICE_DIR)
# 设置离线模式(本机无法访问 huggingface.co # 设置离线模式(本机无法访问 huggingface.co
+4 -4
View File
@@ -6,9 +6,9 @@
原始图 → 人脸检测+1k关键点 → Generator_Matte头发抠图 → 白底合成 + 居中裁剪 → 打标签 原始图 → 人脸检测+1k关键点 → Generator_Matte头发抠图 → 白底合成 + 居中裁剪 → 打标签
用法: 用法:
cd /home/xsl/change_hair/project/hair_service_sd cd /home/ubuntu/change_hair/project/hair_service_sd
/home/xsl/miniconda3/envs/my_hair/bin/python /home/xsl/change_hair/prepare_train_data.py \ /home/ubuntu/miniconda3/envs/my_hair/bin/python /home/ubuntu/change_hair/prepare_train_data.py \
--input /home/xsl/change_hair/train_images \ --input /home/ubuntu/change_hair/train_images \
--hair-id new_hairstyle_001 \ --hair-id new_hairstyle_001 \
--gender boy --gender boy
""" """
@@ -22,7 +22,7 @@ import argparse
ssl._create_default_https_context = ssl._create_unverified_context ssl._create_default_https_context = ssl._create_unverified_context
# 必须在 hair_service_sd 目录运行(依赖其模块) # 必须在 hair_service_sd 目录运行(依赖其模块)
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd" HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
os.chdir(HAIR_SERVICE_DIR) os.chdir(HAIR_SERVICE_DIR)
sys.path.insert(0, HAIR_SERVICE_DIR) sys.path.insert(0, HAIR_SERVICE_DIR)
os.environ.setdefault("HF_HUB_OFFLINE", "1") os.environ.setdefault("HF_HUB_OFFLINE", "1")
+5 -1
View File
@@ -70,7 +70,11 @@ def process_infer(user_img_path, target_color, color_dir, dst_path, ratio=1.0):
# 通过webui对照片进行增强 # 通过webui对照片进行增强
s4 = time.time() s4 = time.time()
# cv2.imwrite('/home/student/Downloads/12121.jpg',crop_img) # cv2.imwrite('/home/student/Downloads/12121.jpg',crop_img)
enhanced_img = enhance_hair.webui_img2img(crop_img, crop_mask, prompt=prompt) try:
enhanced_img = enhance_hair.webui_img2img(crop_img, crop_mask, prompt=prompt)
except Exception as e:
print(f"[process_infer] webui enhance skipped: {e}")
enhanced_img = crop_img
print("----------------------------webui_img2img", time.time() - s4) print("----------------------------webui_img2img", time.time() - s4)
# 写回原图 # 写回原图
+15 -16
View File
@@ -1,21 +1,21 @@
[default] [default]
modelDir = weights modelDir = weights
hairstyleDir = /home/xsl/change_hair/project/data/ref_hairstyle hairstyleDir = /home/ubuntu/change_hair/project/data/ref_hairstyle
haircolorDir= /home/xsl/change_hair/project/data/ref_haircolor haircolorDir= /home/ubuntu/change_hair/project/data/ref_haircolor
userDir=/home/xsl/change_hair/project/data/userImage userDir=/home/ubuntu/change_hair/project/data/userImage
tmp_dir=/home/xsl/change_hair/project/data/tmp tmp_dir=/home/ubuntu/change_hair/project/data/tmp
res_dir=/home/xsl/change_hair/project/data/res_dir res_dir=/home/ubuntu/change_hair/project/data/res_dir
userInfo_dir=/home/xsl/change_hair/project/data/user_info userInfo_dir=/home/ubuntu/change_hair/project/data/user_info
baseColor_ID=HDR10_443322 baseColor_ID=HDR10_443322
Port = 11023 Port = 11023
refer_dir = /home/xsl/change_hair/project/data/ref_online refer_dir = /home/ubuntu/change_hair/project/data/ref_online
ref_user_dir = /home/xsl/change_hair/project/data/ref_user_imgs ref_user_dir = /home/ubuntu/change_hair/project/data/ref_user_imgs
train_dir = /home/xsl/change_hair/data/train_material train_dir = /home/ubuntu/change_hair/data/train_material
refImgDir=/home/xsl/change_hair/project/data/refImage refImgDir=/home/ubuntu/change_hair/project/data/refImage
hair_template_material_dir=/home/xsl/change_hair/project/data/hair_template_material hair_template_material_dir=/home/ubuntu/change_hair/project/data/hair_template_material
ref_color=/home/xsl/change_hair/project/data/ref_color ref_color=/home/ubuntu/change_hair/project/data/ref_color
ref_color_img=/home/xsl/change_hair/project/data/ref_color_imgs ref_color_img=/home/ubuntu/change_hair/project/data/ref_color_imgs
upload_train_dir=/home/xsl/change_hair/project/data/upload_train_imgs upload_train_dir=/home/ubuntu/change_hair/project/data/upload_train_imgs
;version=local ;version=local
version=online version=online
@@ -23,7 +23,7 @@ version=online
strength=1 strength=1
[logger] [logger]
logpath = /home/xsl/change_hair/project/logs logpath = /home/ubuntu/change_hair/project/logs
level=INFO level=INFO
[timelogger] [timelogger]
@@ -32,4 +32,3 @@ name=watch-time
[errorlogger] [errorlogger]
level=ERROR level=ERROR
name=w-error name=w-error
+106 -39
View File
@@ -90,6 +90,14 @@ class HairStyle_Model(object):
self.keypoints_processor = hair_init.keypoints_processor self.keypoints_processor = hair_init.keypoints_processor
self.human_keypoint = hair_init.human_keypoint self.human_keypoint = hair_init.human_keypoint
# 用户图预处理产物内存缓存(一级缓存,避免重复 GPU 推理 + 磁盘 IO)。
# key = 图片字节哈希 + ratiovalue = 7 个产物 + landmarks_1k 的 dict。
# 接口2 女性多发型场景:同一张用户图连续请求,第二次起命中缓存省 ~1.6s。
import collections
self._user_prepare_cache = {} # {cache_key: {产物dict}}
self._user_prepare_cache_keys = collections.deque() # LRU 顺序
self._USER_CACHE_MAX = 8 # 最多缓存 8 张图(防止内存膨胀)
# worker_id = int(os.environ.get('APP_WORKER_ID', 1)) # worker_id = int(os.environ.get('APP_WORKER_ID', 1))
# rand_max = 9527 # rand_max = 9527
@@ -2268,48 +2276,99 @@ class HairStyle_Model(object):
ref_landmark_f1k2_768) ref_landmark_f1k2_768)
cv2.imwrite(another_pose_hair_image_dir, another_pose_hair_image*255) cv2.imwrite(another_pose_hair_image_dir, another_pose_hair_image*255)
landmark1k_dir = osp.join(userinfo_dir, 'kpt_1k.txt') landmark1k_dir = osp.join(userinfo_dir, 'kpt_1k.txt')
if not osp.exists(landmark1k_dir):
landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward_diy(user_rgb_8uc3_orisize) # ===== 内存缓存(一级):key = 图片哈希 + ratio =====
# 命中则跳过 landmark 检测 + get_prepare_user_768_data 整条 GPU 管线(省 ~1.6s
import hashlib as _hashlib
_img_hash = _hashlib.md5(user_rgb_8uc3_orisize.tobytes()).hexdigest()[:16]
_cache_key = f"{_img_hash}_r{ratio}"
_cached = self._user_prepare_cache.get(_cache_key)
if _cached is not None:
# 命中内存缓存:直接取所有产物,跳过 landmark 检测和预处理管线
landmarks_origin_img_1k = _cached['landmarks_1k']
user_bald_res_8uc3_orisize = _cached['bald_res']
user_baldseg_8uc3_orisize = _cached['baldseg_ori']
user_baldseg_8uc3_768 = _cached['baldseg_768']
user_bald_8uc3_768 = _cached['bald_768']
user_landmark_f1k2_768 = _cached['lmk_768']
user_hairstyle_M = _cached['M']
user_matting_8uc3_bald_orisize = _cached['matting_ori']
# 补写磁盘文件:下游代码(功能7等)仍从 userinfo_dir 读这些文件,
# 而 task_id 每次不同导致 userinfo_dir 不同,必须补写保证下游可用
os.makedirs(userinfo_dir, exist_ok=True)
np.savetxt(landmark1k_dir, landmarks_origin_img_1k)
cv2.imwrite(osp.join(userinfo_dir, 'bald_res_ori.png'), user_bald_res_8uc3_orisize)
cv2.imwrite(osp.join(userinfo_dir, 'bald_seg_ori.png'), user_baldseg_8uc3_orisize)
cv2.imwrite(osp.join(userinfo_dir, 'user_baldseg_768.png'), user_baldseg_8uc3_768)
cv2.imwrite(osp.join(userinfo_dir, 'bald_seg_768.png'), user_bald_8uc3_768)
np.savetxt(osp.join(userinfo_dir, 'landmark_f1k2_768.txt'), user_landmark_f1k2_768)
np.savetxt(osp.join(userinfo_dir, 'hairstyle_M.txt'), user_hairstyle_M)
if user_matting_8uc3_bald_orisize is not None:
cv2.imwrite(osp.join(userinfo_dir, 'user_orig_mask.png'), user_matting_8uc3_bald_orisize)
self.logger_process.info(f"内存缓存命中 key={_cache_key},跳过用户图预处理(补写磁盘文件)")
else:
# 未命中:走原有逻辑(landmark 检测 + 预处理管线 + 磁盘缓存)
if not osp.exists(landmark1k_dir):
landmarks_origin_img_1k, bounding_box, euler_info = self.get_landmark.forward_diy(user_rgb_8uc3_orisize)
if landmarks_origin_img_1k is None:
return None, 10001, None, None, None
np.savetxt(landmark1k_dir, landmarks_origin_img_1k)
else:
landmarks_origin_img_1k = np.loadtxt(landmark1k_dir)
# landmarks_origin_img_1k, _, _ = self.get_landmark.forward(user_rgb_8uc3_orisize)
if landmarks_origin_img_1k is None: if landmarks_origin_img_1k is None:
return None, 10001, None, None, None return None, 10001, None, None, None
np.savetxt(landmark1k_dir, landmarks_origin_img_1k)
else:
landmarks_origin_img_1k = np.loadtxt(landmark1k_dir)
# landmarks_origin_img_1k, _, _ = self.get_landmark.forward(user_rgb_8uc3_orisize)
if landmarks_origin_img_1k is None:
return None, 10001, None, None, None
user_bald_res_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_res_ori.png') user_bald_res_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_res_ori.png')
user_baldseg_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_seg_ori.png') user_baldseg_8uc3_orisize_dir = osp.join(userinfo_dir, 'bald_seg_ori.png')
user_baldseg_8uc3_768_dir = osp.join(userinfo_dir, 'user_baldseg_768.png') user_baldseg_8uc3_768_dir = osp.join(userinfo_dir, 'user_baldseg_768.png')
user_bald_8uc3_768_dir = osp.join(userinfo_dir, 'bald_seg_768.png') user_bald_8uc3_768_dir = osp.join(userinfo_dir, 'bald_seg_768.png')
user_landmark_f1k2_768_dir = osp.join(userinfo_dir, 'landmark_f1k2_768.txt') user_landmark_f1k2_768_dir = osp.join(userinfo_dir, 'landmark_f1k2_768.txt')
user_hairstyle_M_dir = osp.join(userinfo_dir, 'hairstyle_M.txt') user_hairstyle_M_dir = osp.join(userinfo_dir, 'hairstyle_M.txt')
condition_exist2 = True condition_exist2 = True
pre_list = [user_bald_res_8uc3_orisize_dir, user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_768_dir, user_bald_8uc3_768_dir, pre_list = [user_bald_res_8uc3_orisize_dir, user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_768_dir, user_bald_8uc3_768_dir,
user_landmark_f1k2_768_dir, user_hairstyle_M_dir] user_landmark_f1k2_768_dir, user_hairstyle_M_dir]
for tmp_dir in pre_list: for tmp_dir in pre_list:
if not osp.exists(tmp_dir): if not osp.exists(tmp_dir):
condition_exist2 = False condition_exist2 = False
user_matting_8uc3_bald_orisize = None user_matting_8uc3_bald_orisize = None
if not condition_exist2: if not condition_exist2:
user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, user_baldseg_8uc3_768, user_bald_8uc3_768, \ user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize, user_baldseg_8uc3_768, user_bald_8uc3_768, \
user_landmark_f1k2_768, user_hairstyle_M, user_matting_8uc3_bald_orisize = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio) user_landmark_f1k2_768, user_hairstyle_M, user_matting_8uc3_bald_orisize = self.process_data.get_prepare_user_768_data(user_rgb_8uc3_orisize, landmarks_origin_img_1k, ratio=ratio)
cv2.imwrite(user_bald_res_8uc3_orisize_dir, user_bald_res_8uc3_orisize) cv2.imwrite(user_bald_res_8uc3_orisize_dir, user_bald_res_8uc3_orisize)
cv2.imwrite(user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_orisize) cv2.imwrite(user_baldseg_8uc3_orisize_dir, user_baldseg_8uc3_orisize)
cv2.imwrite(user_baldseg_8uc3_768_dir, user_baldseg_8uc3_768) cv2.imwrite(user_baldseg_8uc3_768_dir, user_baldseg_8uc3_768)
cv2.imwrite(user_bald_8uc3_768_dir, user_bald_8uc3_768) cv2.imwrite(user_bald_8uc3_768_dir, user_bald_8uc3_768)
np.savetxt(user_landmark_f1k2_768_dir, user_landmark_f1k2_768) np.savetxt(user_landmark_f1k2_768_dir, user_landmark_f1k2_768)
np.savetxt(user_hairstyle_M_dir, user_hairstyle_M) np.savetxt(user_hairstyle_M_dir, user_hairstyle_M)
else: else:
user_bald_res_8uc3_orisize = cv2.imread(user_bald_res_8uc3_orisize_dir) user_bald_res_8uc3_orisize = cv2.imread(user_bald_res_8uc3_orisize_dir)
user_baldseg_8uc3_orisize = cv2.imread(user_baldseg_8uc3_orisize_dir) user_baldseg_8uc3_orisize = cv2.imread(user_baldseg_8uc3_orisize_dir)
user_baldseg_8uc3_768 = cv2.imread(user_baldseg_8uc3_768_dir) user_baldseg_8uc3_768 = cv2.imread(user_baldseg_8uc3_768_dir)
user_bald_8uc3_768 = cv2.imread(user_bald_8uc3_768_dir) user_bald_8uc3_768 = cv2.imread(user_bald_8uc3_768_dir)
user_landmark_f1k2_768 = np.loadtxt(user_landmark_f1k2_768_dir) user_landmark_f1k2_768 = np.loadtxt(user_landmark_f1k2_768_dir)
user_hairstyle_M = np.loadtxt(user_hairstyle_M_dir) user_hairstyle_M = np.loadtxt(user_hairstyle_M_dir)
# 写入内存缓存(含 landmarks_1k 和 matting,供后续同图请求命中)
self._user_prepare_cache[_cache_key] = {
'landmarks_1k': landmarks_origin_img_1k,
'bald_res': user_bald_res_8uc3_orisize,
'baldseg_ori': user_baldseg_8uc3_orisize,
'baldseg_768': user_baldseg_8uc3_768,
'bald_768': user_bald_8uc3_768,
'lmk_768': user_landmark_f1k2_768,
'M': user_hairstyle_M,
'matting_ori': user_matting_8uc3_bald_orisize,
}
self._user_prepare_cache_keys.append(_cache_key)
# LRU 淘汰:超过上限删除最老的
while len(self._user_prepare_cache_keys) > self._USER_CACHE_MAX:
_old = self._user_prepare_cache_keys.popleft()
self._user_prepare_cache.pop(_old, None)
self.logger_process.info(f"内存缓存写入 key={_cache_key},当前缓存 {len(self._user_prepare_cache)}")
# show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1) # show_concat = np.concatenate((user_rgb_8uc3_orisize, user_bald_res_8uc3_orisize, user_baldseg_8uc3_orisize), axis=1)
# ratio = 1536. / max(show_concat.shape[:2]) # ratio = 1536. / max(show_concat.shape[:2])
@@ -2318,18 +2377,26 @@ class HairStyle_Model(object):
# cv2.waitKey() # cv2.waitKey()
user_orig_mask_path = os.path.join(userinfo_dir, "user_orig_mask.png") user_orig_mask_path = os.path.join(userinfo_dir, "user_orig_mask.png")
if not os.path.exists(user_orig_mask_path): if user_matting_8uc3_bald_orisize is not None and not os.path.exists(user_orig_mask_path):
cv2.imwrite(user_orig_mask_path, user_matting_8uc3_bald_orisize) cv2.imwrite(user_orig_mask_path, user_matting_8uc3_bald_orisize)
# 换发型 # 换发型(主 GAN
import time as _time
_t_gan0 = _time.perf_counter()
hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image, hair_gene_8uc3_768 = self.generator_hair.Generator_Hair_inference_use_pref(another_pose_hair_image,
user_baldseg_8uc3_768, user_baldseg_8uc3_768,
user_bald_8uc3_768, user_bald_8uc3_768,
user_landmark_f1k2_768, gender) user_landmark_f1k2_768, gender)
_t_gan = _time.perf_counter() - _t_gan0
# 融合(第3次matte + 融合GAN
_t_fuse0 = _time.perf_counter()
hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste( hair_gene_fusion_8uc3_orisize, hair_gene_matte_8uc3_orisize = self.process_data.get_fusion_res_hairpaste(
user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M) user_bald_res_8uc3_orisize, hair_gene_8uc3_768, user_landmark_f1k2_768, user_hairstyle_M)
_t_fuse = _time.perf_counter() - _t_fuse0
self.logger_process.info(
f"功能6分步计时: 主GAN={_t_gan:.3f}s 融合={_t_fuse:.3f}s (缓存={'命中' if _cached is not None else '未命中'})")
gen_hair_mask_path_2 = os.path.join(userinfo_dir, "hair_mask_2.png") gen_hair_mask_path_2 = os.path.join(userinfo_dir, "hair_mask_2.png")
cv2.imwrite(gen_hair_mask_path_2, hair_gene_matte_8uc3_orisize) cv2.imwrite(gen_hair_mask_path_2, hair_gene_matte_8uc3_orisize)
+5 -4
View File
@@ -1,4 +1,5 @@
import io import io
import os
import os.path import os.path
import time import time
@@ -83,13 +84,13 @@ class ControlnetRequestImg2Img:
return encoded_image return encoded_image
def build_body_v2(self, dst_width, dst_height, cfg_scale, base_img, denoising_strength=0.7): def build_body_v2(self, dst_width, dst_height, cfg_scale, base_img, denoising_strength=0.7, steps=None):
self.body = { self.body = {
"prompt": self.prompt, "prompt": self.prompt,
"negative_prompt": self.neg_prompt, "negative_prompt": self.neg_prompt,
"sampler_name": "DPM++ 2M Karras", "sampler_name": "DPM++ 2M Karras",
"batch_size": 1, "batch_size": 1,
"steps": 20, "steps": int(steps) if steps is not None else int(os.environ.get("WEBUI_STEPS", "15")),
"width": dst_width, "width": dst_width,
"height": dst_height, "height": dst_height,
"cfg_scale": cfg_scale, "cfg_scale": cfg_scale,
@@ -413,7 +414,7 @@ def encode_numpy_to_base64(img):
encoded_image = base64.b64encode(bytes).decode('utf-8') encoded_image = base64.b64encode(bytes).decode('utf-8')
return encoded_image return encoded_image
def webui_img2img(img=None, mask_img=None, in_gender=None, task_id=None, hair_id=None, lora_material_path=None, tag="", is_hr=False, denoising_strength=0.7, inference_port="57860", refiner_switch_at=0.5): def webui_img2img(img=None, mask_img=None, in_gender=None, task_id=None, hair_id=None, lora_material_path=None, tag="", is_hr=False, denoising_strength=0.7, inference_port="57860", refiner_switch_at=0.5, webui_steps=None):
# url = "http://hairservice.tslead.net:57860/sdapi/v1/img2img" # url = "http://hairservice.tslead.net:57860/sdapi/v1/img2img"
neg_prompt = '(nsfw:1.5), ng_deepnegative_v1_75t, (badhandv4:1.2), (worst quality:2), (low quality:2), (normal quality:2), lowres, bad anatomy, bad hands, ((monochrome)), ((grayscale)) watermark, moles, large breast, big breast, bad_pictures,easynegative' neg_prompt = '(nsfw:1.5), ng_deepnegative_v1_75t, (badhandv4:1.2), (worst quality:2), (low quality:2), (normal quality:2), lowres, bad anatomy, bad hands, ((monochrome)), ((grayscale)) watermark, moles, large breast, big breast, bad_pictures,easynegative'
@@ -429,7 +430,7 @@ def webui_img2img(img=None, mask_img=None, in_gender=None, task_id=None, hair_id
control_net = ControlnetRequestImg2Img(prompt, neg_prompt, mask_img) control_net = ControlnetRequestImg2Img(prompt, neg_prompt, mask_img)
# control_net.build_body_hr(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength) # control_net.build_body_hr(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength)
if not is_hr: if not is_hr:
control_net.build_body_v2(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength) control_net.build_body_v2(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength, steps=webui_steps)
else: else:
control_net.build_body_hr(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength, refiner_switch_at=refiner_switch_at) control_net.build_body_hr(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=7, base_img=encoded_image, denoising_strength=denoising_strength, refiner_switch_at=refiner_switch_at)
@@ -72,6 +72,8 @@ class HairStyle_Model_Infer(object):
# haircolor_dir = '/home/data/hair/data/ref_color/3628746832766' # haircolor_dir = '/home/data/hair/data/ref_color/3628746832766'
face_base, hair_matting, status = self.infer_haircolor_new(user_rgb_8uc3_orisize, haircolor_dir,target_hair_color, face_base, hair_matting, status = self.infer_haircolor_new(user_rgb_8uc3_orisize, haircolor_dir,target_hair_color,
return_matting=True) return_matting=True)
if status != 0:
return user_rgb_8uc3_orisize, None, status
user_rgb_8uc3_orisize = face_base user_rgb_8uc3_orisize = face_base
# 构建一个色板 # 构建一个色板
r, g, b = target_hair_color r, g, b = target_hair_color
@@ -1016,7 +1018,7 @@ class HairStyle_Model_Infer(object):
def infer_haircolor_new(self, user_rgb_8uc3_orisize, haircolor_dir, target_hair_color, return_matting=False): def infer_haircolor_new(self, user_rgb_8uc3_orisize, haircolor_dir, target_hair_color, return_matting=False):
landmarks_origin_img_1k= self.get_landmark.forward(user_rgb_8uc3_orisize) landmarks_origin_img_1k= self.get_landmark.forward(user_rgb_8uc3_orisize)
if landmarks_origin_img_1k is None: if landmarks_origin_img_1k is None:
return None, 10001 return None, None, 10001
need_process = (target_hair_color[0] * 0.299 + target_hair_color[1] * 0.587 + target_hair_color[2] * 0.114) > 150 need_process = (target_hair_color[0] * 0.299 + target_hair_color[1] * 0.587 + target_hair_color[2] * 0.114) > 150
_, user_matting_8uc1_bald_orisize = self.process_data_infer.generator_matte.matte_inference(user_rgb_8uc3_orisize, _, user_matting_8uc1_bald_orisize = self.process_data_infer.generator_matte.matte_inference(user_rgb_8uc3_orisize,
landmarks_origin_img_1k) landmarks_origin_img_1k)
@@ -1036,10 +1038,13 @@ class HairStyle_Model_Infer(object):
need_process=True need_process=True
if need_process: if need_process:
haircolor_dir_tmp = os.path.join(config.get('default', "haircolorDir"), config.get('default', "baseColor_ID")) haircolor_dir_tmp = os.path.join(config.get('default', "haircolorDir"), config.get('default', "baseColor_ID"))
face_base, status1 = self.infer_haircolor_tj(user_rgb_8uc3_orisize, haircolor_dir_tmp) if os.path.exists(haircolor_dir_tmp):
# cv2.imwrite('/home/student/Desktop/tmp_color/need/face_base.png', face_base) face_base, status1 = self.infer_haircolor_tj(user_rgb_8uc3_orisize, haircolor_dir_tmp)
if status1 == 0: # cv2.imwrite('/home/student/Desktop/tmp_color/need/face_base.png', face_base)
return face_base, user_matting_8uc1_bald_orisize, 0 if status1 == 0:
return face_base, user_matting_8uc1_bald_orisize, 0
else:
need_process = False
else: else:
need_process = False need_process = False
if not need_process: if not need_process:
@@ -255,7 +255,8 @@ def change_hair_colorv3():
jsonify({'msg': '算法解析错误', 'result': '', 'umd': '', 'state': -1}), 400) jsonify({'msg': '算法解析错误', 'result': '', 'umd': '', 'state': -1}), 400)
except Exception as e: except Exception as e:
print(e) print(f"[hairColor ERROR] {e}")
traceback.print_exc()
return make_response( return make_response(
jsonify({'msg': '算法解析错误', 'result': '', 'umd': '', 'state': -1}), 400) jsonify({'msg': '算法解析错误', 'result': '', 'umd': '', 'state': -1}), 400)
@@ -605,6 +606,8 @@ def change_hairstyle_v4():
else: else:
p_tag = "" p_tag = ""
denoising_strength = float(input_info.get('denoising_strength', 0.6)) # 接口11 可调;默认 0.6 denoising_strength = float(input_info.get('denoising_strength', 0.6)) # 接口11 可调;默认 0.6
webui_steps_raw = input_info.get('webui_steps', None) # 可选:webui img2img 步数,None用服务端默认
webui_steps = int(webui_steps_raw) if webui_steps_raw is not None else None
print(f"功能8:处理发型区域,耗时:{time.time() - start_time:.3f}s") print(f"功能8:处理发型区域,耗时:{time.time() - start_time:.3f}s")
# cv2.imwrite(f"{task_id}_mask_dilate.jpg", mask_dilate) # cv2.imwrite(f"{task_id}_mask_dilate.jpg", mask_dilate)
# cv2.imwrite(f"{task_id}_final_img.jpg", final_img) # cv2.imwrite(f"{task_id}_final_img.jpg", final_img)
@@ -618,7 +621,7 @@ def change_hairstyle_v4():
# mask_dilate = cv2.imread("/root/project/hair_service_sd/gt_mask_dilate.jpg") # mask_dilate = cv2.imread("/root/project/hair_service_sd/gt_mask_dilate.jpg")
sd_result = webui_img2img(img=final_img, mask_img=mask_dilate, in_gender=in_gender, task_id=task_id, sd_result = webui_img2img(img=final_img, mask_img=mask_dilate, in_gender=in_gender, task_id=task_id,
hair_id=hair_id, lora_material_path=hair_material_dir, tag=p_tag, is_hr=is_hr, hair_id=hair_id, lora_material_path=hair_material_dir, tag=p_tag, is_hr=is_hr,
denoising_strength=denoising_strength, inference_port="57860") denoising_strength=denoising_strength, inference_port="57860", webui_steps=webui_steps)
print(f"功能:webui,耗时:{time.time() - start_time:.3f}s") print(f"功能:webui,耗时:{time.time() - start_time:.3f}s")
# 功能:后处理并上传结果 # 功能:后处理并上传结果
@@ -43,6 +43,9 @@ def webui_img2img(img, mask, prompt='', denoising_strength=0.35):
} }
response = requests.post(url=url, json=request_dict) response = requests.post(url=url, json=request_dict)
ret_json = response.json() ret_json = response.json()
if 'images' not in ret_json:
print(f"[webui_img2img ERROR] {ret_json}")
raise Exception(f"webui error: {ret_json.get('error', 'unknown')}")
result = ret_json['images'][0] result = ret_json['images'][0]
img = cv2.imdecode(np.frombuffer(base64.b64decode(result.split(",", 1)[0]), np.uint8), cv2.IMREAD_COLOR) img = cv2.imdecode(np.frombuffer(base64.b64decode(result.split(",", 1)[0]), np.uint8), cv2.IMREAD_COLOR)
return img return img
@@ -37,8 +37,8 @@ if version == "local":
callback_url = 'http://service.aicloud.fit:7395/api/hair/trainCallBack' callback_url = 'http://service.aicloud.fit:7395/api/hair/trainCallBack'
else: else:
current_url = 'http://0.0.0.0:7393/' current_url = 'http://0.0.0.0:7393/'
kohya_ss_home_dir = '/home/xsl/change_hair/project/kohya_ss_home' kohya_ss_home_dir = '/home/ubuntu/change_hair/project/kohya_ss_home'
webui_lora_dir = '/home/xsl/change_hair/project/onediff/stable-diffusion-webui/models/Lora' webui_lora_dir = '/home/ubuntu/change_hair/project/onediff/stable-diffusion-webui/models/Lora'
inference_use_onediff = False inference_use_onediff = False
callback_url = 'http://0.0.0.0:8801/api/hair/trainCallBack' callback_url = 'http://0.0.0.0:8801/api/hair/trainCallBack'
base_webui_port = '57860' base_webui_port = '57860'
@@ -307,13 +307,13 @@ def train_thread(sq, gpu_id):
# 3. GPU 固定 device=0(单卡) # 3. GPU 固定 device=0(单卡)
# 4. 去掉 tokenizer_cache_dir(改用 HF 本地缓存 + 离线模式) # 4. 去掉 tokenizer_cache_dir(改用 HF 本地缓存 + 离线模式)
# 5. 设置 HF_HUB_OFFLINE 避免联网检查 # 5. 设置 HF_HUB_OFFLINE 避免联网检查
kohya_python = '/home/xsl/miniconda3/envs/kohya/bin/python' kohya_python = '/home/ubuntu/miniconda3/envs/my_hair/bin/python'
kohya_workdir = os.path.join(kohya_ss_home_dir, 'kohya_ss') kohya_workdir = os.path.join(kohya_ss_home_dir, 'kohya_ss')
base_model = '/home/xsl/change_hair/project/onediff/stable-diffusion-webui/models/Stable-diffusion/v1-5-pruned-emaonly.safetensors' base_model = '/home/ubuntu/change_hair/project/onediff/stable-diffusion-webui/models/Stable-diffusion/v1-5-pruned-emaonly.safetensors'
cmd_train = ( cmd_train = (
f'cd {kohya_workdir} && ' f'cd {kohya_workdir} && '
f'HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 CUDA_VISIBLE_DEVICES={device_id} ' f'HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 CUDA_VISIBLE_DEVICES={device_id} '
f'/home/xsl/miniconda3/envs/kohya/bin/accelerate launch --num_cpu_threads_per_process=2 "./train_network.py" --enable_bucket ' f'/home/ubuntu/miniconda3/envs/my_hair/bin/accelerate launch --num_cpu_threads_per_process=2 "./train_network.py" --enable_bucket '
f'--min_bucket_reso=256 --max_bucket_reso=2048 --pretrained_model_name_or_path="{base_model}" ' f'--min_bucket_reso=256 --max_bucket_reso=2048 --pretrained_model_name_or_path="{base_model}" '
f'--train_data_dir={images_dir} --resolution="2000,2000" ' f'--train_data_dir={images_dir} --resolution="2000,2000" '
f'--output_dir={model_dir} ' f'--output_dir={model_dir} '
+4 -4
View File
@@ -10,16 +10,16 @@
# bash start_all.sh status # 查看状态 # bash start_all.sh status # 查看状态
# ===================================================================== # =====================================================================
set -u set -u
BASE="/home/xsl/change_hair" BASE="/home/ubuntu/change_hair"
PROJ="$BASE/project" PROJ="$BASE/project"
LOGDIR="$PROJ/logs" LOGDIR="$PROJ/logs"
PIDD="$LOGDIR/pids" PIDD="$LOGDIR/pids"
mkdir -p "$LOGDIR" "$PIDD" mkdir -p "$LOGDIR" "$PIDD"
# conda 环境的 python 路径 # conda 环境的 python 路径
PY_HAIR="/home/xsl/miniconda3/envs/my_hair/bin/python" PY_HAIR="/home/ubuntu/miniconda3/envs/my_hair/bin/python"
PY_SD="/home/xsl/miniconda3/envs/sdwebui/bin/python" PY_SD="/home/ubuntu/miniconda3/envs/sdwebui/bin/python"
PY_PHOTO="/home/xsl/miniconda3/envs/py310/bin/python" PY_PHOTO="/home/ubuntu/miniconda3/envs/py310/bin/python"
# 公共环境变量 # 公共环境变量
export CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1 # 旧版 cryptography 兼容 export CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1 # 旧版 cryptography 兼容
+2 -2
View File
@@ -2,5 +2,5 @@
export CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1 export CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1
export CUDA_VISIBLE_DEVICES=0 export CUDA_VISIBLE_DEVICES=0
export APP_WORKER_ID=1 export APP_WORKER_ID=1
cd /home/xsl/change_hair/project/hair_service_sd cd /home/ubuntu/change_hair/project/hair_service_sd
exec /home/xsl/miniconda3/envs/my_hair/bin/python run_copy_cost_colorb64.py exec /home/ubuntu/miniconda3/envs/my_hair/bin/python run_copy_cost_colorb64.py
+2 -2
View File
@@ -5,5 +5,5 @@ export CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1
export CUDA_VISIBLE_DEVICES=0 export CUDA_VISIBLE_DEVICES=0
export HF_HUB_OFFLINE=1 export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1 export TRANSFORMERS_OFFLINE=1
cd /home/xsl/change_hair/hair_grow_service cd /home/ubuntu/change_hair/hair_grow_service
exec /home/xsl/miniconda3/envs/my_hair/bin/python app.py exec /home/ubuntu/miniconda3/envs/my_hair/bin/python app.py
+2 -2
View File
@@ -1,5 +1,5 @@
#!/bin/bash #!/bin/bash
# 独立启动 photo_service,确保脱离会话 # 独立启动 photo_service,确保脱离会话
export CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1 export CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1
cd /home/xsl/change_hair/project/photo_service cd /home/ubuntu/change_hair/project/photo_service
exec /home/xsl/miniconda3/envs/py310/bin/python -u lora_train_service_1.py exec /home/ubuntu/miniconda3/envs/py310/bin/python -u lora_train_service_1.py
+10 -10
View File
@@ -4,7 +4,7 @@
# 已传数据自动续传(--partial) # 已传数据自动续传(--partial)
set -u set -u
LOGDIR="/home/xsl/change_hair/project/logs" LOGDIR="/home/ubuntu/change_hair/project/logs"
DONE_MARKER="$LOGDIR/sync_all.done" DONE_MARKER="$LOGDIR/sync_all.done"
mkdir -p "$LOGDIR" mkdir -p "$LOGDIR"
@@ -14,15 +14,15 @@ mkdir -p "$LOGDIR"
# 任务定义:名称|源|目标 # 任务定义:名称|源|目标
# 优先级排序:换发色关键链路优先 → 换发型 → 训练相关 → 大数据最后 # 优先级排序:换发色关键链路优先 → 换发型 → 训练相关 → 大数据最后
TASKS=( TASKS=(
"photo_service|szlc@192.168.101.63:/home/szlc/project/photo_service/|/home/xsl/change_hair/project/photo_service/" "photo_service|szlc@192.168.101.63:/home/szlc/project/photo_service/|/home/ubuntu/change_hair/project/photo_service/"
"project_data|szlc@192.168.101.63:/home/szlc/project/data/|/home/xsl/change_hair/project/data/" "project_data|szlc@192.168.101.63:/home/szlc/project/data/|/home/ubuntu/change_hair/project/data/"
"hair_service_sd|szlc@192.168.101.63:/home/szlc/project/hair_service_sd/|/home/xsl/change_hair/project/hair_service_sd/" "hair_service_sd|szlc@192.168.101.63:/home/szlc/project/hair_service_sd/|/home/ubuntu/change_hair/project/hair_service_sd/"
"conda_py310|szlc@192.168.101.63:/home/szlc/miniconda3/envs/py310/|/home/xsl/miniconda3/envs/py310/" "conda_py310|szlc@192.168.101.63:/home/szlc/miniconda3/envs/py310/|/home/ubuntu/miniconda3/envs/py310/"
"conda_my_hair|szlc@192.168.101.63:/home/szlc/miniconda3/envs/my_hair/|/home/xsl/miniconda3/envs/my_hair/" "conda_my_hair|szlc@192.168.101.63:/home/szlc/miniconda3/envs/my_hair/|/home/ubuntu/miniconda3/envs/my_hair/"
"conda_sdwebui|szlc@192.168.101.63:/home/szlc/miniconda3/envs/sdwebui/|/home/xsl/miniconda3/envs/sdwebui/" "conda_sdwebui|szlc@192.168.101.63:/home/szlc/miniconda3/envs/sdwebui/|/home/ubuntu/miniconda3/envs/sdwebui/"
"kohya_ss_home|szlc@192.168.101.63:/home/szlc/project/kohya_ss_home/|/home/xsl/change_hair/project/kohya_ss_home/" "kohya_ss_home|szlc@192.168.101.63:/home/szlc/project/kohya_ss_home/|/home/ubuntu/change_hair/project/kohya_ss_home/"
"onediff|szlc@192.168.101.63:/home/szlc/project/onediff/|/home/xsl/change_hair/project/onediff/" "onediff|szlc@192.168.101.63:/home/szlc/project/onediff/|/home/ubuntu/change_hair/project/onediff/"
"train_material|szlc@192.168.101.63:/data/train_material/|/home/xsl/change_hair/data/train_material/" "train_material|szlc@192.168.101.63:/data/train_material/|/home/ubuntu/change_hair/data/train_material/"
) )
echo "========================================" | tee -a "$LOGDIR/sync_main.log" echo "========================================" | tee -a "$LOGDIR/sync_main.log"
+2 -2
View File
@@ -1,8 +1,8 @@
#!/bin/bash #!/bin/bash
# 重新同步 conda 环境(上次被 sed 损坏,这次纯净拷贝) # 重新同步 conda 环境(上次被 sed 损坏,这次纯净拷贝)
set -u set -u
CONDA_ENVS="/home/xsl/miniconda3/envs" CONDA_ENVS="/home/ubuntu/miniconda3/envs"
LOGDIR="/home/xsl/change_hair/project/logs" LOGDIR="/home/ubuntu/change_hair/project/logs"
mkdir -p "$CONDA_ENVS" "$LOGDIR" mkdir -p "$CONDA_ENVS" "$LOGDIR"
echo "conda 环境重传开始: $(date '+%H:%M:%S')" | tee -a "$LOGDIR/sync_conda.log" echo "conda 环境重传开始: $(date '+%H:%M:%S')" | tee -a "$LOGDIR/sync_conda.log"
+3 -3
View File
@@ -2,8 +2,8 @@
"""端到端测试:生发走换发型工作流""" """端到端测试:生发走换发型工作流"""
import base64, requests, time, sys import base64, requests, time, sys
img_path = "/home/xsl/change_hair/project/data/userImage/8488902485_20250630055547.jpg" img_path = "/home/ubuntu/change_hair/project/data/userImage/8488902485_20250630055547.jpg"
mask_path = "/home/xsl/change_hair/project/logs/hairgrow_test_mask.png" mask_path = "/home/ubuntu/change_hair/project/logs/hairgrow_test_mask.png"
with open(img_path, "rb") as f: with open(img_path, "rb") as f:
img_b64 = "data:image/jpeg;base64," + base64.b64encode(f.read()).decode() img_b64 = "data:image/jpeg;base64," + base64.b64encode(f.read()).decode()
@@ -24,7 +24,7 @@ try:
d = r.json() d = r.json()
print(f"HTTP {r.status_code} | state={d.get('state')} | msg={d.get('msg','')} | 耗时={time.time()-t0:.1f}s") print(f"HTTP {r.status_code} | state={d.get('state')} | msg={d.get('msg','')} | 耗时={time.time()-t0:.1f}s")
if d.get("state") == 0: if d.get("state") == 0:
out = "/home/xsl/change_hair/project/logs/test_grow_swap_result.jpg" out = "/home/ubuntu/change_hair/project/logs/test_grow_swap_result.jpg"
with open(out, "wb") as f: with open(out, "wb") as f:
f.write(base64.b64decode(d["result"])) f.write(base64.b64decode(d["result"]))
print(f"✅ 生发成功! 结果图: {out}") print(f"✅ 生发成功! 结果图: {out}")
+2 -2
View File
@@ -2,7 +2,7 @@
"""换发色接口测试""" """换发色接口测试"""
import base64, json, requests, sys import base64, json, requests, sys
img_path = "/home/xsl/change_hair/project/data/userImage/8488902485_20250630055547.jpg" img_path = "/home/ubuntu/change_hair/project/data/userImage/8488902485_20250630055547.jpg"
with open(img_path, "rb") as f: with open(img_path, "rb") as f:
img_b64 = base64.b64encode(f.read()).decode() img_b64 = base64.b64encode(f.read()).decode()
@@ -21,7 +21,7 @@ try:
result = d.get("result", "") result = d.get("result", "")
if d.get("state") == 0 and result: if d.get("state") == 0 and result:
# 保存结果图 # 保存结果图
out = "/home/xsl/change_hair/project/logs/test_haircolor_result.jpg" out = "/home/ubuntu/change_hair/project/logs/test_haircolor_result.jpg"
with open(out, "wb") as f: with open(out, "wb") as f:
f.write(base64.b64decode(result)) f.write(base64.b64decode(result))
print(f"✅ 换发色成功! 结果图已保存: {out} ({len(result)} bytes base64)") print(f"✅ 换发色成功! 结果图已保存: {out} ({len(result)} bytes base64)")
+3 -3
View File
@@ -2,8 +2,8 @@
"""测试 /api/hairGrow/v1 接口""" """测试 /api/hairGrow/v1 接口"""
import base64, requests, sys import base64, requests, sys
img_path = "/home/xsl/change_hair/project/data/userImage/8488902485_20250630055547.jpg" img_path = "/home/ubuntu/change_hair/project/data/userImage/8488902485_20250630055547.jpg"
mask_path = "/home/xsl/change_hair/project/logs/hairgrow_test_mask.png" mask_path = "/home/ubuntu/change_hair/project/logs/hairgrow_test_mask.png"
with open(img_path, "rb") as f: with open(img_path, "rb") as f:
img_b64 = "data:image/jpeg;base64," + base64.b64encode(f.read()).decode() img_b64 = "data:image/jpeg;base64," + base64.b64encode(f.read()).decode()
@@ -26,7 +26,7 @@ try:
print(f"HTTP {r.status_code} | msg={d.get('msg')} | state={d.get('state')} | 耗时={time.time()-t0:.1f}s") print(f"HTTP {r.status_code} | msg={d.get('msg')} | state={d.get('state')} | 耗时={time.time()-t0:.1f}s")
result = d.get("result", "") result = d.get("result", "")
if d.get("state") == 0 and result: if d.get("state") == 0 and result:
out = "/home/xsl/change_hair/project/logs/test_hairgrow_api_result.jpg" out = "/home/ubuntu/change_hair/project/logs/test_hairgrow_api_result.jpg"
with open(out, "wb") as f: with open(out, "wb") as f:
f.write(base64.b64decode(result)) f.write(base64.b64decode(result))
print(f"✅ 生发成功! 结果图: {out}") print(f"✅ 生发成功! 结果图: {out}")
+2 -2
View File
@@ -2,7 +2,7 @@
"""换发型接口测试""" """换发型接口测试"""
import base64, json, requests, sys, time import base64, json, requests, sys, time
img_path = "/home/xsl/change_hair/project/data/userImage/8488902485_20250630055547.jpg" img_path = "/home/ubuntu/change_hair/project/data/userImage/8488902485_20250630055547.jpg"
with open(img_path, "rb") as f: with open(img_path, "rb") as f:
img_b64 = base64.b64encode(f.read()).decode() img_b64 = base64.b64encode(f.read()).decode()
@@ -22,7 +22,7 @@ try:
print(f"HTTP {r.status_code} | msg={d.get('msg')} | state={d.get('state')} | 耗时={time.time()-t0:.1f}s") print(f"HTTP {r.status_code} | msg={d.get('msg')} | state={d.get('state')} | 耗时={time.time()-t0:.1f}s")
data = d.get("data", "") data = d.get("data", "")
if d.get("state") == 0 and data: if d.get("state") == 0 and data:
out = "/home/xsl/change_hair/project/logs/test_swaphair_result.jpg" out = "/home/ubuntu/change_hair/project/logs/test_swaphair_result.jpg"
with open(out, "wb") as f: with open(out, "wb") as f:
f.write(base64.b64decode(data)) f.write(base64.b64decode(data))
print(f"✅ 换发型成功! 结果图已保存: {out} ({len(data)} bytes base64)") print(f"✅ 换发型成功! 结果图已保存: {out} ({len(data)} bytes base64)")
+1 -1
View File
@@ -4,7 +4,7 @@ import requests, json
# 用一个已有完整训练素材的发型 # 用一个已有完整训练素材的发型
HAIR_ID = "1905785164224868354" HAIR_ID = "1905785164224868354"
MATERIAL_DIR = f"/home/xsl/change_hair/data/train_material/{HAIR_ID}" MATERIAL_DIR = f"/home/ubuntu/change_hair/data/train_material/{HAIR_ID}"
payload = { payload = {
"task_id": f"test_train_{HAIR_ID}", "task_id": f"test_train_{HAIR_ID}",
+7 -7
View File
@@ -15,7 +15,7 @@ import time
import json import json
import shutil import shutil
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd" HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
os.chdir(HAIR_SERVICE_DIR) os.chdir(HAIR_SERVICE_DIR)
sys.path.insert(0, HAIR_SERVICE_DIR) sys.path.insert(0, HAIR_SERVICE_DIR)
os.environ.setdefault("HF_HUB_OFFLINE", "1") os.environ.setdefault("HF_HUB_OFFLINE", "1")
@@ -26,15 +26,15 @@ import requests as req
# (hair_id, gender, template_img) # (hair_id, gender, template_img)
HAIRSTYLES = [ HAIRSTYLES = [
("chang_tuoyuan", "girl", "/home/xsl/change_hair/hair_type_images/chang_tuoyuan/chang_tuoyuan.jpg"), ("chang_tuoyuan", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_tuoyuan/chang_tuoyuan.jpg"),
("chang_bolang", "girl", "/home/xsl/change_hair/hair_type_images/chang_bolang/chang_bolang.jpg"), ("chang_bolang", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_bolang/chang_bolang.jpg"),
("chang_zhixian", "girl", "/home/xsl/change_hair/hair_type_images/chang_zhixian/chang_zhixian.jpg"), ("chang_zhixian", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_zhixian/chang_zhixian.jpg"),
("chang_huaban", "girl", "/home/xsl/change_hair/hair_type_images/chang_huaban/chang_huaban.jpg"), ("chang_huaban", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_huaban/chang_huaban.jpg"),
("chang_xinxing", "girl", "/home/xsl/change_hair/hair_type_images/chang_xinxing/chang_xinxing.jpg"), ("chang_xinxing", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_xinxing/chang_xinxing.jpg"),
] ]
HAIR_CALLBACK = "http://127.0.0.1:8801/api/hair/trainCallBack" HAIR_CALLBACK = "http://127.0.0.1:8801/api/hair/trainCallBack"
UPLOAD_TRAIN_DIR = "/home/xsl/change_hair/project/data/upload_train_imgs" UPLOAD_TRAIN_DIR = "/home/ubuntu/change_hair/project/data/upload_train_imgs"
# 导入 step4/step5 函数(会触发模型加载) # 导入 step4/step5 函数(会触发模型加载)
from train_hairstyle_full import step4_template_material, step5_preview from train_hairstyle_full import step4_template_material, step5_preview
+10 -6
View File
@@ -6,8 +6,8 @@
串接 prepare_train_data / gen_template_material / photo_service训练 / 回调 / 预览图 串接 prepare_train_data / gen_template_material / photo_service训练 / 回调 / 预览图
用法: 用法:
cd /home/xsl/change_hair/project/hair_service_sd cd /home/ubuntu/change_hair/project/hair_service_sd
python /home/xsl/change_hair/train_hairstyle_full.py --hair-id huaban1 --input /home/xsl/change_hair/huaban1 --gender girl --template-img /home/xsl/change_hair/huaban1/huaban1.jpg python /home/ubuntu/change_hair/train_hairstyle_full.py --hair-id huaban1 --input /home/ubuntu/change_hair/huaban1 --gender girl --template-img /home/ubuntu/change_hair/huaban1/huaban1.jpg
""" """
import os import os
import sys import sys
@@ -21,7 +21,7 @@ import argparse
import requests as req import requests as req
ssl._create_default_https_context = ssl._create_unverified_context ssl._create_default_https_context = ssl._create_unverified_context
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd" HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
os.chdir(HAIR_SERVICE_DIR) os.chdir(HAIR_SERVICE_DIR)
sys.path.insert(0, HAIR_SERVICE_DIR) sys.path.insert(0, HAIR_SERVICE_DIR)
os.environ.setdefault("HF_HUB_OFFLINE", "1") os.environ.setdefault("HF_HUB_OFFLINE", "1")
@@ -48,9 +48,9 @@ RESOLUTIONS = [512, 768, 1024, 1280, 1536]
PHOTO_TRAIN = "http://127.0.0.1:32678/api/hair/train" PHOTO_TRAIN = "http://127.0.0.1:32678/api/hair/train"
HAIR_CALLBACK = "http://127.0.0.1:8801/api/hair/trainCallBack" HAIR_CALLBACK = "http://127.0.0.1:8801/api/hair/trainCallBack"
SWAP_API = "http://127.0.0.1:8801/api/swapHair/v1" SWAP_API = "http://127.0.0.1:8801/api/swapHair/v1"
GIRL_IMG = "/home/xsl/change_hair/images/girl.png" GIRL_IMG = "/home/ubuntu/change_hair/images/girl.png"
BOY_IMG = "/home/xsl/change_hair/images/boy.png" BOY_IMG = "/home/ubuntu/change_hair/images/boy.png"
PREVIEW_DIR = "/home/xsl/change_hair/hair_grow_service/static/previews" PREVIEW_DIR = "/home/ubuntu/change_hair/hair_grow_service/static/previews"
_detector = _aligner = _matte = None _detector = _aligner = _matte = None
def get_models(): def get_models():
@@ -158,6 +158,10 @@ def step3_wait_and_callback(hair_id, template_img):
else: else:
print(" ✗ 训练超时"); return False print(" ✗ 训练超时"); return False
time.sleep(5) # 等训练进程写完 time.sleep(5) # 等训练进程写完
# 训练时停掉了 hair 服务,回调前需要重启
print(" 重启 change_hair-hair 服务...")
os.system("sudo systemctl restart change_hair-hair")
time.sleep(20)
# 触发回调生成ref材质 # 触发回调生成ref材质
print(" 触发回调生成ref材质...") print(" 触发回调生成ref材质...")
r = req.post(HAIR_CALLBACK, json={"task_id":f"train_{hair_id}","hair_id":hair_id, r = req.post(HAIR_CALLBACK, json={"task_id":f"train_{hair_id}","hair_id":hair_id,
+9 -9
View File
@@ -8,8 +8,8 @@
阶段C: 启动服务后串行 step3(回调生成ref材质) + step4(模板) + step5(预览) 阶段C: 启动服务后串行 step3(回调生成ref材质) + step4(模板) + step5(预览)
用法: 用法:
cd /home/xsl/change_hair/project/hair_service_sd cd /home/ubuntu/change_hair/project/hair_service_sd
python /home/xsl/change_hair/train_hairstyles_parallel.py python /home/ubuntu/change_hair/train_hairstyles_parallel.py
""" """
import os import os
import sys import sys
@@ -18,7 +18,7 @@ import json
import subprocess import subprocess
from concurrent.futures import ThreadPoolExecutor, as_completed from concurrent.futures import ThreadPoolExecutor, as_completed
HAIR_SERVICE_DIR = "/home/xsl/change_hair/project/hair_service_sd" HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
os.chdir(HAIR_SERVICE_DIR) os.chdir(HAIR_SERVICE_DIR)
sys.path.insert(0, HAIR_SERVICE_DIR) sys.path.insert(0, HAIR_SERVICE_DIR)
os.environ.setdefault("HF_HUB_OFFLINE", "1") os.environ.setdefault("HF_HUB_OFFLINE", "1")
@@ -29,11 +29,11 @@ os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")
# ====== 配置:本次要训练的发型 ====== # ====== 配置:本次要训练的发型 ======
# (hair_id, input_dir, gender, template_img) # (hair_id, input_dir, gender, template_img)
HAIRSTYLES = [ HAIRSTYLES = [
("chang_tuoyuan", "/home/xsl/change_hair/hair_type_images/chang_tuoyuan", "girl", "/home/xsl/change_hair/hair_type_images/chang_tuoyuan/chang_tuoyuan.jpg"), ("chang_tuoyuan", "/home/ubuntu/change_hair/hair_type_images/chang_tuoyuan", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_tuoyuan/chang_tuoyuan.jpg"),
("chang_bolang", "/home/xsl/change_hair/hair_type_images/chang_bolang", "girl", "/home/xsl/change_hair/hair_type_images/chang_bolang/chang_bolang.jpg"), ("chang_bolang", "/home/ubuntu/change_hair/hair_type_images/chang_bolang", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_bolang/chang_bolang.jpg"),
("chang_zhixian", "/home/xsl/change_hair/hair_type_images/chang_zhixian", "girl", "/home/xsl/change_hair/hair_type_images/chang_zhixian/chang_zhixian.jpg"), ("chang_zhixian", "/home/ubuntu/change_hair/hair_type_images/chang_zhixian", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_zhixian/chang_zhixian.jpg"),
("chang_huaban", "/home/xsl/change_hair/hair_type_images/chang_huaban", "girl", "/home/xsl/change_hair/hair_type_images/chang_huaban/chang_huaban.jpg"), ("chang_huaban", "/home/ubuntu/change_hair/hair_type_images/chang_huaban", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_huaban/chang_huaban.jpg"),
("chang_xinxing", "/home/xsl/change_hair/hair_type_images/chang_xinxing", "girl", "/home/xsl/change_hair/hair_type_images/chang_xinxing/chang_xinxing.jpg"), ("chang_xinxing", "/home/ubuntu/change_hair/hair_type_images/chang_xinxing", "girl", "/home/ubuntu/change_hair/hair_type_images/chang_xinxing/chang_xinxing.jpg"),
] ]
PARALLEL = 2 # LoRA 训练并发数 PARALLEL = 2 # LoRA 训练并发数
@@ -137,7 +137,7 @@ def main():
log(f"阶段B完成: LoRA 训练完成 {len(completed)}/{len(step1_ok)}") log(f"阶段B完成: LoRA 训练完成 {len(completed)}/{len(step1_ok)}")
# 写一个完成清单文件,供阶段C脚本读取 # 写一个完成清单文件,供阶段C脚本读取
done_file = "/home/xsl/change_hair/project/logs/train_batch_done.json" done_file = "/home/ubuntu/change_hair/project/logs/train_batch_done.json"
with open(done_file, "w") as f: with open(done_file, "w") as f:
json.dump({"completed": sorted(completed), "all": ids}, f) json.dump({"completed": sorted(completed), "all": ids}, f)
log(f"已写入完成清单: {done_file}") log(f"已写入完成清单: {done_file}")
+6 -6
View File
@@ -6,7 +6,7 @@
通过 ThreadPoolExecutor 控制并发数 通过 ThreadPoolExecutor 控制并发数
用法: 用法:
python /home/xsl/change_hair/train_lora_parallel.py python /home/ubuntu/change_hair/train_lora_parallel.py
前提step1 训练数据已准备好data/train_material/<hid>/images/1_hairstyle/*.png 前提step1 训练数据已准备好data/train_material/<hid>/images/1_hairstyle/*.png
""" """
@@ -16,11 +16,11 @@ import time
import subprocess import subprocess
from concurrent.futures import ThreadPoolExecutor, as_completed from concurrent.futures import ThreadPoolExecutor, as_completed
TRAIN_DIR = "/home/xsl/change_hair/data/train_material" TRAIN_DIR = "/home/ubuntu/change_hair/data/train_material"
KOHYA_WORKDIR = "/home/xsl/change_hair/project/kohya_ss_home/kohya_ss" KOHYA_WORKDIR = "/home/ubuntu/change_hair/project/kohya_ss_home/kohya_ss"
BASE_MODEL = "/home/xsl/change_hair/project/onediff/stable-diffusion-webui/models/Stable-diffusion/v1-5-pruned-emaonly.safetensors" BASE_MODEL = "/home/ubuntu/change_hair/project/onediff/stable-diffusion-webui/models/Stable-diffusion/v1-5-pruned-emaonly.safetensors"
KOHYA_ACCEL = "/home/xsl/miniconda3/envs/kohya/bin/accelerate" KOHYA_ACCEL = "/home/ubuntu/miniconda3/envs/kohya/bin/accelerate"
LOG_DIR = "/home/xsl/change_hair/project/logs" LOG_DIR = "/home/ubuntu/change_hair/project/logs"
# 本次要训练的发型 # 本次要训练的发型
HAIRSTYLES = ["chang_tuoyuan", "chang_zhixian", "chang_huaban", "chang_xinxing"] HAIRSTYLES = ["chang_tuoyuan", "chang_zhixian", "chang_huaban", "chang_xinxing"]