初始化：换发型/换发色/训练发型服务
包含： - 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密钥已脱敏为环境变量，原文件备份在本地
This commit is contained in:
xsl
@@ -0,0 +1,190 @@
|
||||
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
|
||||
|
||||
|
||||
# 将当前工作目录切换到当前目录
|
||||
project_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
os.chdir(project_dir)
|
||||
sys.path.append(project_dir)
|
||||
|
||||
import json
|
||||
|
||||
def send_request(url, state, task_dict):
|
||||
import requests
|
||||
import json
|
||||
|
||||
msg = '头发lora训练失败' if state == -1 else '头发lora训练成功'
|
||||
payload = json.dumps({
|
||||
"task_id": task_dict['task_id'] if task_dict is not None else '',
|
||||
"hair_id": task_dict['hair_id'] if task_dict is not None else '',
|
||||
"state": state,
|
||||
"msg": msg
|
||||
})
|
||||
headers = {
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
|
||||
requests.request("POST", url, headers=headers, data=payload)
|
||||
|
||||
|
||||
|
||||
def train_hair_lora():
|
||||
import os
|
||||
import json
|
||||
task_id = ''
|
||||
try:
|
||||
# 获取请求参数
|
||||
request_data = ""
|
||||
assert 'task_id' in request_data, 'task_id is required'
|
||||
task_id = request_data['task_id']
|
||||
assert 'hair_id' in request_data and 'hair_material_dir' in request_data, 'hair_id and hair_material_dir is required'
|
||||
hair_material_dir = request_data['hair_material_dir']
|
||||
assert os.path.exists(hair_material_dir), f'{hair_material_dir} not exists'
|
||||
img_dir = os.path.join(hair_material_dir,'images')
|
||||
assert os.path.exists(img_dir), f'{img_dir} not exists'
|
||||
assert os.path.exists(os.path.join(hair_material_dir, 'model')), f'{os.path.join(hair_material_dir, "model")} not exists'
|
||||
|
||||
#判断img_dir下面是否只有一个文件夹
|
||||
img_dir_list = os.listdir(img_dir)
|
||||
assert len(img_dir_list) == 1, f'{img_dir}下面只能有一个文件夹'
|
||||
train_image_dir = os.path.join(img_dir,img_dir_list[0])
|
||||
assert os.path.isdir(train_image_dir), f'{train_image_dir}不是文件夹'
|
||||
# 检查该文件夹是否以数字加下划线开头
|
||||
assert img_dir_list[0].split('_')[0].isdigit(), f'{img_dir_list[0]}不是数字开头'
|
||||
|
||||
#判断文件夹下面是否有图片
|
||||
img_list = glob.glob(train_image_dir + '/*.png')
|
||||
assert len(img_list) > 10, f'{os.path.join(img_dir,img_dir_list[0])}下面图片数量小于10张'
|
||||
request_data['train_image_dir'] = train_image_dir
|
||||
except Exception as e:
|
||||
print(e)
|
||||
|
||||
|
||||
def inference_webui():
|
||||
import os
|
||||
import requests
|
||||
url = "http://127.0.0.1:57860/sdapi/v1/img2img"
|
||||
try:
|
||||
# 获取请求参数
|
||||
request_data = ""
|
||||
# with open('request_data.json', 'w') as f:
|
||||
# json.dump(request_data, f)
|
||||
request_json = request_data.get('request_json')
|
||||
hair_material_dir = request_data.get('hair_material_dir')
|
||||
assert os.path.exists(hair_material_dir), f'{hair_material_dir} not exists'
|
||||
request_lora_path = os.path.join(hair_material_dir, 'model', 'hairstyle_lora.safetensors')
|
||||
assert os.path.exists(request_lora_path), f'{request_lora_path} not exists'
|
||||
|
||||
tmp_lora_name = f'{uuid.uuid4()}'
|
||||
lora_dst_path = ""
|
||||
os.system('cp {} {}'.format(request_lora_path, lora_dst_path))
|
||||
request_json['prompt'] = f'<lora:{tmp_lora_name}:1.0>, titor hairstyle, ' + request_json['prompt']
|
||||
request_json['negative_prompt'] = request_json['negative_prompt'] + ', faceless, no human'
|
||||
response = requests.post(url=url, json=request_json)
|
||||
ret_json = response.json()
|
||||
|
||||
if os.path.exists(lora_dst_path):
|
||||
os.remove(lora_dst_path)
|
||||
|
||||
except Exception as e:
|
||||
print(e)
|
||||
|
||||
|
||||
def tag(train_image_dir):
|
||||
img_path_list = glob.glob(train_image_dir + '/*.png')
|
||||
|
||||
# 给训练图片打标签, gpt
|
||||
for img_path in img_path_list:
|
||||
print(img_path)
|
||||
txt_path = img_path[:-4] + '.txt'
|
||||
# if not os.path.exists(txt_path):
|
||||
tags = caption_image(img_path)
|
||||
with open(txt_path, 'w') as f:
|
||||
f.write('titor hairstyle, faceless, no human, gray background, simple background, ' + tags)
|
||||
# else:
|
||||
# continue
|
||||
|
||||
|
||||
def train_thread(sq, gpu_id):
|
||||
while True:
|
||||
url = 'http://service.aicloud.fit:7393/api/hair/trainCallBack'
|
||||
task_dict = None
|
||||
try:
|
||||
task_dict = sq.get()
|
||||
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']
|
||||
# 给训练图片打标签
|
||||
# cmd_caption = ('docker run --rm --gpus all -v /home/chinatszrn/Documents/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 "./finetune/tag_images_by_wd14_tagger.py" --batch_size=2 '
|
||||
# '--general_threshold=0.5 --character_threshold=0.5 --caption_extension=".txt" '
|
||||
# '--model="SmilingWolf/wd-v1-4-convnextv2-tagger-v2" --max_data_loader_n_workers="2" '
|
||||
# '--debug --remove_underscore --frequency_tags --undesired_tags="1girl, 1boy" '
|
||||
# f'"{train_image_dir}"')
|
||||
# os.system(cmd_caption)
|
||||
|
||||
img_path_list = glob.glob(train_image_dir + '/*.png')
|
||||
# 给训练图片打标签, gpt
|
||||
for img_path in img_path_list:
|
||||
tags = caption_image(img_path)
|
||||
txt_path = img_path[:-4]+'.txt'
|
||||
if not os.path.exists(txt_path):
|
||||
with open(txt_path, 'w') as f:
|
||||
f.write('titor hairstyle, faceless, no human, gray background, simple background, ' + tags)
|
||||
else:
|
||||
continue
|
||||
|
||||
|
||||
#判断是否每个训练图片都有标签文件
|
||||
# for img_path in img_path_list:
|
||||
# assert os.path.exists(img_path[:-4]+'.txt'), f'{img_path}没有对应的标签文件'
|
||||
# with open(img_path[:-4]+'.txt', 'r') as f:
|
||||
# tags = f.readline()
|
||||
# with open(img_path[:-4]+'.txt', 'w') as f:
|
||||
# f.write('titor hairstyle, faceless, no human, gray background, simple background, ' + tags)
|
||||
|
||||
#训练头发lora
|
||||
cmd_train = ('docker run --rm --gpus all -v /home/chinatszrn/Documents/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=512 --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="1800,1800" '
|
||||
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="12" '
|
||||
'--no_half_vae --learning_rate="0.0001" --lr_scheduler="cosine" --lr_warmup_steps="96" --train_batch_size="1" '
|
||||
'--max_train_steps="4000" --save_every_n_epochs="100" --mixed_precision="fp16" --save_precision="fp16" '
|
||||
'--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"')
|
||||
os.system(cmd_train)
|
||||
|
||||
lora_path = os.path.join(model_dir, 'hairstyle_lora.safetensors')
|
||||
|
||||
|
||||
if not os.path.exists(lora_path):
|
||||
send_request(url, -1, task_dict)
|
||||
else:
|
||||
send_request(url, 0, task_dict)
|
||||
|
||||
except Exception as e:
|
||||
send_request(url, -1, None)
|
||||
continue
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
tag("/mnt/database2/online-server/hair-online/hair_lora_train_material/AAVWGW1NN-0KC-33B-24-/images/1_hairstyle")
|
||||
Reference in New Issue
Block a user