colomi
0eb61f3e60
包含： - hair_service_sd: 换发型/换发色算法服务 (端口 8801) - photo_service: LoRA 训练调度服务 (端口 32678) - stable-diffusion-webui: SD WebUI 推理服务 (端口 57860) - kohya_ss_home: 训练环境代码 - meidaojia: 监控测试脚本 - setup.sh: 一键部署脚本 (conda环境恢复 + 配置生成 + 完整性检查) - start_all_services.sh: 启动3个服务 - configure.ini.template: 路径模板化 (BASE_DIR自动推导) - conda_envs/py310.yml: py310 环境定义 大文件 (weights/, models/, data/, conda_envs/*.tar.gz 等) 通过 .gitignore 排除， 由网盘单独上传。
160 lines
5.6 KiB
Python
160 lines
5.6 KiB
Python
import os
|
|
import sys
|
|
from webui_im2im import ControlnetRequestImg2Img
|
|
import numpy as np
|
|
import base64
|
|
import cv2
|
|
import os,sys
|
|
from gevent import pywsgi, monkey
|
|
|
|
monkey.patch_all()
|
|
|
|
# 将当前工作目录切换到当前目录
|
|
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
|
|
|
|
app = Flask(__name__)
|
|
|
|
|
|
@app.route('/template/list', methods=['POST'])
|
|
def get_template_list():
|
|
import os
|
|
import json
|
|
try:
|
|
# 读取data/template.json文件，并返回
|
|
with open('data/template.json', 'r') as f:
|
|
template_list = json.load(f)
|
|
ret_dict = dict(code=0, message='success', data=template_list)
|
|
# 返回结果作为 JSON 响应
|
|
return jsonify(ret_dict)
|
|
|
|
except Exception as e:
|
|
ret_dict = dict(code=-1, message=str(e), data=[])
|
|
return jsonify(ret_dict)
|
|
|
|
|
|
@app.route('/user/list', methods=['POST'])
|
|
def get_user_list():
|
|
try:
|
|
# 获取用户文件夹路径
|
|
user_list_dir = os.path.join(global_var.service_data_dir, 'user_data')
|
|
|
|
ret_user_list = []
|
|
|
|
# 遍历查找user_list_dir目录下所有的cfg.json文件
|
|
for root, dirs, files in os.walk(user_list_dir):
|
|
for file in files:
|
|
if file != 'cfg.json': continue
|
|
json_path = os.path.join(root, file)
|
|
lora_path = os.path.join(root, 'lora.safetensors')
|
|
if not os.path.exists(lora_path): continue
|
|
# 读取cfg.json文件
|
|
with open(json_path, 'r') as f:
|
|
user_info = json.load(f)
|
|
user_dict = dict(user_id=user_info['user_id'], face_img_url=user_info['face_img_url'])
|
|
ret_user_list.append(user_dict)
|
|
|
|
ret_dict = dict(code=0, message='success', data=ret_user_list)
|
|
return jsonify(ret_dict)
|
|
|
|
except Exception as e:
|
|
ret_dict = dict(code=-1, message='error', data=[])
|
|
return jsonify(ret_dict)
|
|
|
|
|
|
@app.route('/user/generate', methods=['POST'])
|
|
def generate_photo():
|
|
try:
|
|
# 获取请求参数
|
|
request_data = request.get_json()
|
|
#判断'user_id'和'base_img'是否在请求参数中
|
|
assert 'user_id' in request_data and 'base_img' in request_data, 'user_id and base_img is required'
|
|
user_id = request_data['user_id']
|
|
base_img_b64 = request_data['base_img']
|
|
|
|
user_path = os.path.join(global_var.service_data_dir, 'user_data', user_id)
|
|
assert os.path.isdir(user_path), 'user_id not exist'
|
|
user_lora = os.path.join(user_path, 'lora.safetensors')
|
|
usr_config_path = os.path.join(user_path, 'cfg.json')
|
|
assert os.path.exists(user_lora) and os.path.exists(usr_config_path), 'lora file or config not exist'
|
|
|
|
# 读取用户配置文件
|
|
with open(usr_config_path, 'r') as f:
|
|
user_info = json.load(f)
|
|
lora_md5 = user_info['lora_md5']
|
|
dst_lora_path = os.path.join(global_var.webui_lora_dir, lora_md5 + '.safetensors')
|
|
if not os.path.exists(dst_lora_path):
|
|
os.system(f'cp {user_lora} {dst_lora_path}')
|
|
|
|
#将模板图像转化为numpy数组
|
|
prompt = f'<lora:{lora_md5}:0.8>,easyphoto_face, easyphoto, 1person,face,suit'
|
|
neg_prompt = '(worst quality:2),(low quality:2),(normal quality:2),lowres,watermark'
|
|
|
|
image_array = np.frombuffer(base64.b64decode(base_img_b64), np.uint8)
|
|
base_img = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
|
|
# base_img = cv2.resize(base_img, (512, 512))
|
|
# cv2.imshow('image', base_img)
|
|
# cv2.waitKey()
|
|
|
|
# 生成图片
|
|
control_net = ControlnetRequestImg2Img(prompt, neg_prompt)
|
|
control_net.build_body(dst_width=base_img.shape[1], dst_height=base_img.shape[0], cfg_scale=3.5, base_img=base_img)
|
|
output = control_net.send_request()
|
|
generate_photo = output['images'][0]
|
|
|
|
# 清理硬盘空间
|
|
os.remove(dst_lora_path)
|
|
|
|
# # 将生成的图片转化为base64编码
|
|
# retval, bytes = cv2.imencode('.png', generate_photo)
|
|
# generate_photo = base64.b64encode(bytes).decode('utf-8')
|
|
return jsonify(dict(code=0, message='success', generate_photo_b64=generate_photo))
|
|
|
|
except Exception as e:
|
|
ret_dict = dict(code=-1, message=str(e), data=[])
|
|
return jsonify(ret_dict)
|
|
|
|
|
|
def webd_service():
|
|
# 用于启动webd的后台服务
|
|
current_file_dir = os.path.dirname(os.path.abspath(__file__))
|
|
webd_path = os.path.join(current_file_dir, 'webd', 'webd')
|
|
|
|
print('webd server started...')
|
|
cmd = f"{webd_path} -w {global_var.service_data_dir} -g rlT -l 10219"
|
|
os.system(cmd)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# 服务启动的数据目录
|
|
global_var.service_data_dir = sys.argv[1]
|
|
|
|
# 本地webui的lora存储目录
|
|
global_var.webui_lora_dir = sys.argv[2]
|
|
|
|
# webui_server_port
|
|
global_var.webui_server_port = int(sys.argv[3])
|
|
|
|
# server_port
|
|
global_var.server_port = int(sys.argv[4])
|
|
|
|
# 检查webui_lora_dir目录是否存在
|
|
assert os.path.isdir(global_var.webui_lora_dir), 'webui_lora_dir should be a directory'
|
|
|
|
# 检查service_data_dir目录是否存在
|
|
if not os.path.exists(global_var.service_data_dir):
|
|
os.makedirs(global_var.service_data_dir, exist_ok=True)
|
|
else:
|
|
assert os.path.isdir(global_var.service_data_dir), 'service_data_dir should be a directory'
|
|
|
|
# 启动服务
|
|
# app.run(debug=False, port=global_var.server_port, host='0.0.0.0')
|
|
|
|
server = pywsgi.WSGIServer(('0.0.0.0', global_var.server_port), app) # test port
|
|
server.serve_forever()
|