colomi
40a512a6c4
- lora_train_service_1.py: 训练底模改为 v1-5-pruned-emaonly.safetensors - gen_super_image.py: 注释掉 refiner_checkpoint（majicmix 不存在会导致推理报错） - setup.sh: 检查项改为 v1-5-pruned-emaonly.safetensors - README.md: 移除 majicmixRealistic_v7 下载步骤
358 lines
15 KiB
Python
358 lines
15 KiB
Python
# -*- coding: utf-8 -*-
|
|
|
|
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
|
|
import datetime
|
|
# 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)
|
|
|
|
from flask import Flask, request, jsonify
|
|
import global_variable as global_var
|
|
import json
|
|
#from common.logger import config
|
|
|
|
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
version = "online"
|
|
if version == "local":
|
|
current_url = 'http://service.aicloud.fit:7393/'
|
|
kohya_ss_home_dir = '/home/chinatszrn/Documents/miaoya/kohya_ss_home'
|
|
webui_lora_dir = '/home/chinatszrn/Documents/miaoya/webui_home/stable-diffusion-webui/models/Lora'
|
|
inference_use_onediff = False
|
|
callback_url = 'http://service.aicloud.fit:7395/api/hair/trainCallBack'
|
|
else:
|
|
current_url = 'http://0.0.0.0:7393/'
|
|
kohya_ss_home_dir = os.path.join(BASE_DIR, 'kohya_ss_home')
|
|
webui_lora_dir = os.path.join(BASE_DIR, 'stable-diffusion-webui', 'models', 'Lora')
|
|
inference_use_onediff = False
|
|
callback_url = 'http://0.0.0.0:8801/api/hair/trainCallBack'
|
|
base_webui_port = '57860'
|
|
base_onediff_port = '9038'
|
|
base_webui_url = "http://127.0.0.1:57860/sdapi/v1/img2img"
|
|
|
|
|
|
app = Flask(__name__)
|
|
|
|
|
|
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,
|
|
"is_tj": task_dict['is_tj']
|
|
})
|
|
|
|
print("payload:", payload)
|
|
|
|
headers = {
|
|
'Content-Type': 'application/json'
|
|
}
|
|
|
|
requests.request("POST", url, headers=headers, data=payload)
|
|
|
|
|
|
@app.route('/api/hair/train', methods=['POST'])
|
|
def train_hair_lora():
|
|
import os
|
|
import json
|
|
task_id = ''
|
|
try:
|
|
# 获取请求参数
|
|
request_data = request.get_json()
|
|
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
|
|
|
|
# is tianjin company
|
|
try:
|
|
request_data['is_tj'] = request_data['is_tj']
|
|
request_data['device_id'] = request_data['device_id']
|
|
except Exception as e:
|
|
print(e)
|
|
request_data['is_tj'] = '0'
|
|
request_data['device_id'] = '1'
|
|
|
|
# 将请求数据放入队列
|
|
global_var.hair_style_lora_train_task_sq.put(request_data)
|
|
|
|
# 返回结果
|
|
ret_dict = dict(state=0, msg='头发lora训练开始', task_id=task_id)
|
|
return jsonify(ret_dict)
|
|
except Exception as e:
|
|
ret_dict = dict(state=-1, msg=str(e), task_id=task_id)
|
|
return jsonify(ret_dict)
|
|
|
|
|
|
@app.route('/api/hair/inference', methods=['POST'])
|
|
def inference_webui():
|
|
import os
|
|
import requests
|
|
import time
|
|
|
|
|
|
try:
|
|
t1 = time.time()
|
|
# 获取请求参数
|
|
request_data = request.get_json()
|
|
# with open('request_data.json', 'w') as f:
|
|
# json.dump(request_data, f)
|
|
request_json = request_data.get('request_json')
|
|
# hd_version_flag = request_data.get('hd_version_flag', False)
|
|
hair_id = request_data['hair_id']
|
|
inference_port = request_data['inference_port']
|
|
|
|
if inference_use_onediff and 'refiner_checkpoint' not in request_json:
|
|
selected_url = base_webui_url.replace(base_webui_port, base_onediff_port)
|
|
request_json['script_name'] = 'onediff_diffusion_model'
|
|
print('user onediff')
|
|
else:
|
|
selected_url = base_webui_url.replace(base_webui_port, inference_port)
|
|
|
|
|
|
|
|
# if hd_version_flag:
|
|
# print('user hd version!!!!!')
|
|
# lora_file_name = 'hairstyle_hd_lora.safetensors'
|
|
# else:
|
|
# lora_file_name = 'hairstyle_lora.safetensors'
|
|
|
|
hair_material_dir = request_data.get('hair_material_dir')
|
|
print(hair_material_dir)
|
|
assert os.path.exists(hair_material_dir), f'{hair_material_dir} not exists'
|
|
|
|
hd_version_flag = True
|
|
lora_file_name = 'hairstyle_hd_lora.safetensors'
|
|
request_lora_path = os.path.join(hair_material_dir, 'model', lora_file_name)
|
|
if not os.path.exists(request_lora_path):
|
|
print('use low resolution lora')
|
|
hd_version_flag = False
|
|
lora_file_name = 'hairstyle_lora.safetensors'
|
|
request_lora_path = os.path.join(hair_material_dir, 'model', lora_file_name)
|
|
|
|
assert os.path.exists(request_lora_path), f'{request_lora_path} not exists'
|
|
|
|
if hd_version_flag:
|
|
tmp_lora_name = f'{hair_id}_hd'
|
|
else:
|
|
tmp_lora_name = f'{hair_id}'
|
|
lora_dst_path = os.path.join(global_var.webui_lora_dir, f'{tmp_lora_name}.safetensors')
|
|
if not os.path.exists(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'
|
|
print('pre_stage cost:', time.time() - t1)
|
|
|
|
start_time = time.time()
|
|
current_time = datetime.datetime.now()
|
|
print("*********cxt log1, infer***************, {}:{}:{}, {}".format(current_time.hour, current_time.minute,
|
|
current_time.second, selected_url))
|
|
response = requests.post(url=selected_url, json=request_json)
|
|
print('inference time:', time.time() - start_time)
|
|
|
|
ret_json = response.json()
|
|
# if os.path.exists(lora_dst_path):
|
|
# os.remove(lora_dst_path)
|
|
# 返回结果
|
|
return jsonify(ret_json)
|
|
except Exception as e:
|
|
ret_dict = dict(state=-1, msg=str(e))
|
|
return jsonify(ret_dict)
|
|
|
|
|
|
@app.route('/api/hair/inference_diy', methods=['POST'])
|
|
def inference_diy_webui():
|
|
import os
|
|
import requests
|
|
import time
|
|
|
|
try:
|
|
t1 = time.time()
|
|
# 获取请求参数
|
|
request_data = request.get_json()
|
|
# with open('request_data.json', 'w') as f:
|
|
# json.dump(request_data, f)
|
|
request_json = request_data.get('request_json')
|
|
inference_port = request_data['inference_port']
|
|
|
|
if inference_use_onediff and 'refiner_checkpoint' not in request_json:
|
|
selected_url = base_webui_url.replace(base_webui_port, base_onediff_port)
|
|
request_json['script_name'] = 'onediff_diffusion_model'
|
|
print('user onediff')
|
|
else:
|
|
selected_url = base_webui_url.replace(base_webui_port, inference_port)
|
|
|
|
request_json['prompt'] = f'titor hairstyle, ' + request_json['prompt']
|
|
request_json['negative_prompt'] = request_json['negative_prompt'] + ', faceless, no human'
|
|
print('pre_stage cost:', time.time() - t1)
|
|
|
|
start_time = time.time()
|
|
current_time = datetime.datetime.now()
|
|
print("*********cxt log1, diy***************, {}:{}:{}, {}".format(current_time.hour, current_time.minute,
|
|
current_time.second, selected_url))
|
|
response = requests.post(url=selected_url, json=request_json)
|
|
print('inference time:', time.time() - start_time)
|
|
|
|
ret_json = response.json()
|
|
# if os.path.exists(lora_dst_path):
|
|
# os.remove(lora_dst_path)
|
|
# 返回结果
|
|
return jsonify(ret_json)
|
|
except Exception as e:
|
|
ret_dict = dict(state=-1, msg=str(e))
|
|
return jsonify(ret_dict)
|
|
|
|
def train_thread(sq, gpu_id):
|
|
while True:
|
|
url = f'{current_url}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']
|
|
# tag = task_dict['tag']
|
|
tag = ""
|
|
is_tj = task_dict['is_tj']
|
|
device_id = task_dict['device_id']
|
|
|
|
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 768 '
|
|
'--w 768 --s 30 --l 7')
|
|
|
|
# 给训练图片打标签
|
|
# 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:
|
|
# print(img_path)
|
|
txt_path = img_path[:-4] + '.txt'
|
|
|
|
print("tag img_path: ", img_path)
|
|
# tags = caption_image(img_path)
|
|
tags = tag
|
|
with open(txt_path, 'w') as f:
|
|
f.write('titor hairstyle, easyphoto, faceless, no human, white background, simple background, ' + tags)
|
|
|
|
#判断是否每个训练图片都有标签文件
|
|
# 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)
|
|
|
|
os.system(f'chmod 777 {model_dir}')
|
|
#训练头发lora
|
|
sd_models_dir = os.path.join(BASE_DIR, 'stable-diffusion-webui', 'models', 'Stable-diffusion')
|
|
container_images_dir = images_dir.replace(kohya_ss_home_dir, '/home/chinatszrn')
|
|
container_model_dir = model_dir.replace(kohya_ss_home_dir, '/home/chinatszrn')
|
|
cmd_train = (
|
|
f'sudo docker run --rm --privileged=true --gpus "device=0" '
|
|
f'-v {kohya_ss_home_dir}:/home/chinatszrn '
|
|
f'-v {sd_models_dir}:/mnt/sd_models '
|
|
'-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=2048 --pretrained_model_name_or_path="/mnt/sd_models/v1-5-pruned-emaonly.safetensors" '
|
|
f'--train_data_dir={container_images_dir} --resolution="2000,2000" '
|
|
f'--output_dir={container_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_hd_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="1500" --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="1000" '
|
|
'--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_hd_lora.safetensors')
|
|
|
|
# tianjin callback
|
|
if is_tj == "1":
|
|
url = f'{callback_url}'
|
|
|
|
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__':
|
|
global_var.webui_lora_dir = webui_lora_dir
|
|
|
|
# avatar训练的任务队列
|
|
global_var.hair_style_lora_train_task_sq = Queue()
|
|
p = Process(target=train_thread, args=(
|
|
global_var.hair_style_lora_train_task_sq, 0))
|
|
p.start()
|
|
|
|
# # 启动服务
|
|
# app.run(debug=True, port=32678, host='0.0.0.0')
|
|
|
|
server = pywsgi.WSGIServer(('0.0.0.0', 32678), app) # test port
|
|
server.serve_forever()
|
|
|
|
p.join()
|