xsl
443cfa298f
包含： - hair_service_sd: 主服务（换发型/换发色/生发，端口8801） - photo_service: LoRA调度+训练（端口32678） - hair_grow_service: 调试测试页（端口8888，含4个测试页） - 批量训练脚本（batch_train_hairstyles.py） - 发际线mask自动识别（hairline_mask.py，4种方案） - 手绘mask换发型（hair_swap_manual.py） - 文档：README.md + LARGE_FILES.md + docs/ 大文件（模型权重200G、训练数据123G）已排除，见 LARGE_FILES.md OSS/COS密钥已脱敏为环境变量，原文件备份在本地
121 lines
4.5 KiB
Python
121 lines
4.5 KiB
Python
import os
|
|
import sys
|
|
import uuid
|
|
|
|
from webui_im2im import ControlnetRequestImg2Img
|
|
import numpy as np
|
|
import base64
|
|
import cv2
|
|
import os, sys
|
|
from gevent import pywsgi, monkey
|
|
from multiprocessing import Process, Queue
|
|
import glob
|
|
|
|
# from gpt4v_caption import caption_image
|
|
|
|
# monkey.patch_all()
|
|
# sys.setrecursionlimit(20000)
|
|
|
|
# 将当前工作目录切换到当前目录
|
|
project_dir = os.path.dirname(os.path.abspath(__file__))
|
|
os.chdir(project_dir)
|
|
sys.path.append(project_dir)
|
|
|
|
|
|
import global_variable as global_var
|
|
import json
|
|
|
|
|
|
|
|
def train_hair_lora():
|
|
import os
|
|
import json
|
|
|
|
try:
|
|
hair_train_dir = "/mnt/database2/jiangqian/0808/online_train_datas_2"
|
|
hair_material_dir_list = os.listdir(hair_train_dir)
|
|
|
|
for single_hair_material in hair_material_dir_list:
|
|
# 获取请求参数
|
|
task_id = str(uuid.uuid4())
|
|
|
|
request_data = {}
|
|
hair_material_dir = os.path.join(hair_train_dir, single_hair_material)
|
|
print("--------------------hair_material_dir:", hair_material_dir)
|
|
request_data['hair_material_dir'] = hair_material_dir
|
|
|
|
img_dir = os.path.join(hair_material_dir, 'images')
|
|
|
|
model_dir = os.path.join(hair_material_dir, 'model')
|
|
if not os.path.exists(model_dir):
|
|
os.makedirs(model_dir)
|
|
|
|
#判断img_dir下面是否只有一个文件夹
|
|
img_dir_list = os.listdir(img_dir)
|
|
train_image_dir = os.path.join(img_dir, img_dir_list[0])
|
|
|
|
#判断文件夹下面是否有图片
|
|
request_data['train_image_dir'] = train_image_dir
|
|
|
|
# 将请求数据放入队列
|
|
train_thread(request_data)
|
|
|
|
# 返回结果
|
|
print("头发lora训练开始")
|
|
print("\n\n\n\n")
|
|
|
|
except Exception as e:
|
|
print(e)
|
|
|
|
def train_thread(task_dict):
|
|
try:
|
|
hair_material_dir = task_dict['hair_material_dir']
|
|
images_dir = os.path.join(hair_material_dir, 'images')
|
|
model_dir = os.path.join(hair_material_dir, 'model')
|
|
train_image_dir = task_dict['train_image_dir']
|
|
tag = ""
|
|
|
|
sample_dir = os.path.join(model_dir, 'sample')
|
|
if not os.path.exists(sample_dir):
|
|
os.makedirs(sample_dir)
|
|
sample_txt = os.path.join(sample_dir, 'prompt.txt')
|
|
with open(sample_txt, 'w') as f_s:
|
|
f_s.write('titor hairstyle, easyphoto, faceless, no human, white background, simple background, '
|
|
+ tag
|
|
+ ' --n low quality, worst quality, bad anatomy, bad composition, poor, low effort --h 512 '
|
|
'--w 512 --s 30 --l 7')
|
|
|
|
#训练头发lora
|
|
cmd_train = (
|
|
'docker run --rm --gpus all -v /home/student/Documents/workspace_cxt_tianjing_hair/miaoya/kohya_ss_home:/home/chinatszrn -v '
|
|
'/mnt:/mnt -e PATH=/home/chinatszrn/.local/bin -w /home/chinatszrn/kohya_ss '
|
|
'--net=host chinatszrn/ubuntu:kohya_ss accelerate launch --num_cpu_threads_per_process=2 "./train_network.py" --enable_bucket '
|
|
'--min_bucket_reso=256 --max_bucket_reso=1800 --pretrained_model_name_or_path="/mnt/nas_hdd/米亚像馆/models/Stable-diffusion/majicmixRealistic_v7.safetensors" '
|
|
f'--train_data_dir={images_dir} --resolution="768,768" '
|
|
f'--output_dir={model_dir} '
|
|
'--network_alpha="64" --save_model_as=safetensors --network_module=networks.lora --text_encoder_lr=5e-05 '
|
|
'--unet_lr=0.0001 --network_dim=128 --output_name="hairstyle_lora" --lr_scheduler_num_cycles="20" '
|
|
'--no_half_vae --learning_rate="0.0001" --lr_scheduler="cosine" --lr_warmup_steps="650" --train_batch_size="1" '
|
|
'--max_train_steps="2000" --save_every_n_epochs="100" --mixed_precision="fp16" --save_precision="fp16" '
|
|
'--caption_extension=".txt" --sample_sampler=ddim '
|
|
f'--sample_prompts={sample_txt} --sample_every_n_epochs="1" '
|
|
'--seed="1234" --cache_latents --optimizer_type="AdamW8bit" --max_data_loader_n_workers="0" --bucket_reso_steps=64 '
|
|
'--xformers --bucket_no_upscale --noise_offset=0.0 --tokenizer_cache_dir="/home/chinatszrn/.cache/clip"')
|
|
print("cmd_train:", cmd_train)
|
|
os.system(cmd_train)
|
|
|
|
lora_path = os.path.join(model_dir, 'hairstyle_lora.safetensors')
|
|
|
|
except Exception as e:
|
|
print(e)
|
|
return
|
|
|
|
|
|
if __name__ == '__main__':
|
|
global_var.webui_lora_dir = '/home/student/Documents/workspace_cxt_tianjing_hair/miaoya/webui_home/stable-diffusion-webui/models/Lora'
|
|
|
|
train_hair_lora()
|
|
|
|
|
|
|