# -*- 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' # kohya 训练环境路径（本地执行，不再使用 Docker） CONDA_BASE = os.environ.get('CONDA_BASE', os.path.expanduser('~/miniconda3')) kohya_python = os.path.join(CONDA_BASE, 'envs', 'kohya', 'bin', 'python') kohya_local_packages = os.path.join(kohya_ss_home_dir, '.local', 'lib', 'python3.10', 'site-packages') kohya_ss_code_dir = os.path.join(kohya_ss_home_dir, 'kohya_ss') sd_models_dir = os.path.join(BASE_DIR, 'stable-diffusion-webui', 'models', 'Stable-diffusion') clip_cache_dir = os.path.join(kohya_ss_home_dir, '.cache', 'clip') 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', 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（本地执行，使用 kohya conda 环境 + .local 包） base_model_path = os.path.join(sd_models_dir, 'v1-5-pruned-emaonly.safetensors') cmd_train = ( f'cd "{kohya_ss_code_dir}" && ' f'PYTHONPATH="{kohya_local_packages}" ' f'"{kohya_python}" -m accelerate.commands.launch --num_cpu_threads_per_process=2 "./train_network.py" --enable_bucket ' '--min_bucket_reso=256 --max_bucket_reso=2048 ' f'--pretrained_model_name_or_path="{base_model_path}" ' f'--train_data_dir="{images_dir}" --resolution="2000,2000" ' 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_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 ' f'--xformers --bucket_no_upscale --noise_offset=0.0 --tokenizer_cache_dir="{clip_cache_dir}"') 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()