save code
This commit is contained in:
+556
@@ -0,0 +1,556 @@
|
||||
import os
|
||||
import uuid
|
||||
import requests
|
||||
from flask import Flask, request, jsonify
|
||||
from urllib.parse import urlparse
|
||||
import json
|
||||
import base64
|
||||
import time
|
||||
from PIL import Image
|
||||
from datetime import datetime
|
||||
from volcenginesdkarkruntime import Ark
|
||||
import shutil
|
||||
import oss2
|
||||
import redis
|
||||
from check_img_body import is_thigh_visible
|
||||
import logging
|
||||
from config import QUEUE_NAME, DEFAULT_TIMEOUT, KEY_QUEUE_LOCK_NAME, acquire_lock
|
||||
|
||||
|
||||
APP_ROOT = os.path.dirname(os.path.abspath(__file__))
|
||||
STATIC_FOLDER = os.path.join(APP_ROOT, 'static')
|
||||
|
||||
# 创建 logger
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.setLevel(logging.INFO) # 设置 logger 的级别
|
||||
|
||||
# 获取当前日期和时间
|
||||
now = datetime.now()
|
||||
|
||||
# 格式化为字符串(例如:2023-10-25 14:30:45)
|
||||
date_time_str = now.strftime("%Y-%m-%d_%H:%M:%S")
|
||||
|
||||
# 创建文件 handler
|
||||
file_handler = logging.FileHandler(f'/var/log/fuyan/change_app_{date_time_str}.log')
|
||||
file_handler.setLevel(logging.INFO)
|
||||
file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
|
||||
file_handler.setFormatter(file_formatter)
|
||||
|
||||
# 创建控制台 handler
|
||||
console_handler = logging.StreamHandler()
|
||||
console_handler.setLevel(logging.INFO)
|
||||
console_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
|
||||
console_handler.setFormatter(console_formatter)
|
||||
|
||||
# 添加 handlers 到 logger
|
||||
logger.addHandler(file_handler)
|
||||
logger.addHandler(console_handler)
|
||||
|
||||
from config import REDIS_HOST
|
||||
from config import REDIS_PORT
|
||||
from config import REDIS_DB
|
||||
|
||||
# 创建 Redis 连接池
|
||||
redis_pool = redis.ConnectionPool(
|
||||
host=REDIS_HOST,
|
||||
port=REDIS_PORT,
|
||||
db=REDIS_DB,
|
||||
max_connections=2000 # 根据实际情况调整
|
||||
)
|
||||
|
||||
def get_redis_conn():
|
||||
"""获取 Redis 连接"""
|
||||
return redis.Redis(connection_pool=redis_pool)
|
||||
|
||||
app = Flask(__name__)
|
||||
|
||||
client = Ark(
|
||||
# 此为默认路径,您可根据业务所在地域进行配置
|
||||
base_url="https://ark.cn-beijing.volces.com/api/v3",
|
||||
# 从环境变量中获取您的 API Key。此为默认方式,您可根据需要进行修改
|
||||
# api_key=os.environ.get("ARK_API_KEY"),
|
||||
api_key='14fc0280-fc65-462d-ac2d-50178c0212e3'
|
||||
)
|
||||
|
||||
def GetPicDesc(img_url):
|
||||
response = client.chat.completions.create(
|
||||
# 指定您创建的方舟推理接入点 ID,此处已帮您修改为您的推理接入点 ID
|
||||
model="doubao-1.5-vision-pro-250328",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": img_url
|
||||
},
|
||||
},
|
||||
{"type": "text", "text": "图片是一件服装的照片,返回为json格式, 第一个字段'服装长度', 图片中的服装穿到正常人身上的长度会覆盖到人体的哪个部位, 只要答案, 选项有(胸、腰、跨、大腿、膝盖、小腿、脚踝、拖地、难以辨认). 第二个字段'衣袖',只要答案,选项有(短袖、长袖)"},
|
||||
],
|
||||
}
|
||||
],
|
||||
|
||||
)
|
||||
|
||||
text = response.choices[0].message.content
|
||||
print(text)
|
||||
return text
|
||||
|
||||
def takeoff_cloth_first(human_name):
|
||||
print('换泳装')
|
||||
queue = requests.get("http://localhost:8188/queue").json()
|
||||
if queue["queue_running"] or queue["queue_pending"]:
|
||||
return None, "cur gpu is busy"
|
||||
|
||||
with open('/home/szlc/code/ComfyUI/change_cloth/change_new.json', 'r', encoding='utf-8') as file:
|
||||
prompt_text = file.read()
|
||||
|
||||
prompt = json.loads(prompt_text)
|
||||
|
||||
prompt["96"]["inputs"]["cloth_len"] = '脚踝'
|
||||
prompt["96"]["inputs"]["cloth_short"] = True
|
||||
|
||||
# prompt["99"]["inputs"]["width"] = int((c_width/c_height) * 1024)
|
||||
|
||||
#input cloth img
|
||||
prompt["22"]["inputs"]["image"] = 'cloth_black.jpg'
|
||||
|
||||
#input human img
|
||||
prompt["61"]["inputs"]["image"] = human_name
|
||||
|
||||
#out put name
|
||||
out_img_name = str(uuid.uuid4())[:8]
|
||||
prompt["102"]["inputs"]["filename_prefix"] = out_img_name
|
||||
|
||||
p = {"prompt": prompt}
|
||||
|
||||
data = json.dumps(p).encode('utf-8')
|
||||
|
||||
response = requests.post("http://localhost:8188/prompt", data=data)
|
||||
prompt_id = response.json()["prompt_id"]
|
||||
|
||||
|
||||
# 2. 轮询队列,直到任务完成
|
||||
while True:
|
||||
queue = requests.get("http://localhost:8188/queue").json()
|
||||
# print(queue)
|
||||
if not queue["queue_running"] and not queue["queue_pending"]:
|
||||
break # 队列为空,任务已完成
|
||||
time.sleep(0.5) # 避免频繁请求
|
||||
|
||||
# 3. 从历史记录中获取结果
|
||||
history = requests.get("http://localhost:8188/history").json()
|
||||
# print("History:", history)
|
||||
outputs = history[prompt_id]["outputs"]
|
||||
for out in outputs:
|
||||
if 'images' in outputs[out]:
|
||||
for out_img in outputs[out]['images']:
|
||||
if 'filename' in out_img:
|
||||
out_img_file_name = out_img['filename']
|
||||
if out_img_name in out_img_file_name:
|
||||
return out_img_file_name
|
||||
return None
|
||||
|
||||
def generate_from_face(human_name, is_girl = True):
|
||||
print('生成写真')
|
||||
queue = requests.get("http://localhost:8188/queue").json()
|
||||
if queue["queue_running"] or queue["queue_pending"]:
|
||||
return None, "cur gpu is busy"
|
||||
|
||||
with open('/home/szlc/code/ComfyUI/change_cloth/xiezhen_girl.json', 'r', encoding='utf-8') as file:
|
||||
prompt_text = file.read()
|
||||
|
||||
prompt = json.loads(prompt_text)
|
||||
|
||||
|
||||
prompt["93"]["inputs"]["image"] = human_name
|
||||
|
||||
if is_girl:
|
||||
prompt["6"]["inputs"]["text"] = "full body shot(1.9),complete figure,front-facing,model pose, Hands are at the sides of the body, smile,clear skin glow,wearing black micro skirt,simple white background,clean composition,32k,high detail,professional beauty photography style,"
|
||||
else:
|
||||
prompt["6"]["inputs"]["text"] = "full body shot(1.9),complete figure,front-facing,model pose,Hands are at the sides of the body,smile,clear skin glow,wearing fitted vest and tight shortst,simple white background,clean composition,32k,high detail,professional beauty photography style,"
|
||||
|
||||
#out put name
|
||||
out_img_name = str(uuid.uuid4())[:8]
|
||||
prompt["65"]["inputs"]["filename_prefix"] = out_img_name
|
||||
|
||||
p = {"prompt": prompt}
|
||||
|
||||
data = json.dumps(p).encode('utf-8')
|
||||
|
||||
response = requests.post("http://localhost:8188/prompt", data=data)
|
||||
prompt_id = response.json()["prompt_id"]
|
||||
|
||||
|
||||
# 2. 轮询队列,直到任务完成
|
||||
while True:
|
||||
queue = requests.get("http://localhost:8188/queue").json()
|
||||
# print(queue)
|
||||
if not queue["queue_running"] and not queue["queue_pending"]:
|
||||
break # 队列为空,任务已完成
|
||||
time.sleep(0.5) # 避免频繁请求
|
||||
|
||||
# 3. 从历史记录中获取结果
|
||||
history = requests.get("http://localhost:8188/history").json()
|
||||
# print("History:", history)
|
||||
outputs = history[prompt_id]["outputs"]
|
||||
for out in outputs:
|
||||
if 'images' in outputs[out]:
|
||||
for out_img in outputs[out]['images']:
|
||||
if 'filename' in out_img:
|
||||
out_img_file_name = out_img['filename']
|
||||
if out_img_name in out_img_file_name:
|
||||
return out_img_file_name
|
||||
return None
|
||||
|
||||
def change(human_name, cloth_name, c_width, c_height, cloth_url, is_generated):
|
||||
json_str = GetPicDesc(cloth_url)
|
||||
json_data = json.loads(json_str)
|
||||
cloth_len = json_data['服装长度']
|
||||
cloth_short = True
|
||||
if '长袖' in json_data['衣袖']:
|
||||
cloth_short = False
|
||||
if cloth_len == '难以辨认':
|
||||
return None, "get image type error"
|
||||
|
||||
thigh_visible = is_thigh_visible(f"/home/szlc/code/ComfyUI/input/{human_name}")
|
||||
if thigh_visible:
|
||||
#脱衣服
|
||||
takeoff_file_name = takeoff_cloth_first(human_name)
|
||||
if takeoff_file_name == None:
|
||||
return None, f"takeoff_cloth_first error {human_name}"
|
||||
else:
|
||||
human_name = takeoff_file_name
|
||||
takeoff_file_path_name = os.path.join('/home/szlc/code/ComfyUI/output', takeoff_file_name)
|
||||
takeoff_file_path_name_input = os.path.join('/home/szlc/code/ComfyUI/input', takeoff_file_name)
|
||||
shutil.copy(takeoff_file_path_name, takeoff_file_path_name_input)
|
||||
else:
|
||||
#生成写真
|
||||
generate_name = generate_from_face(human_name)
|
||||
if generate_name == None:
|
||||
return None, f"takeoff_cloth_first error {human_name}"
|
||||
else:
|
||||
human_name = generate_name
|
||||
generate_name_file_path_name = os.path.join('/home/szlc/code/ComfyUI/output', generate_name)
|
||||
generate_name_file_path_name_input = os.path.join('/home/szlc/code/ComfyUI/input', generate_name)
|
||||
shutil.copy(generate_name_file_path_name, generate_name_file_path_name_input)
|
||||
|
||||
queue = requests.get("http://localhost:8188/queue").json()
|
||||
if queue["queue_running"] or queue["queue_pending"]:
|
||||
return None, "cur gpu is busy"
|
||||
|
||||
print('开始换衣服')
|
||||
with open('/home/szlc/code/ComfyUI/change_cloth/change_new.json', 'r', encoding='utf-8') as file:
|
||||
prompt_text = file.read()
|
||||
|
||||
prompt = json.loads(prompt_text)
|
||||
|
||||
prompt["96"]["inputs"]["cloth_len"] = cloth_len
|
||||
prompt["96"]["inputs"]["cloth_short"] = cloth_short
|
||||
|
||||
# prompt["99"]["inputs"]["width"] = int((c_width/c_height) * 1024)
|
||||
|
||||
#input cloth img
|
||||
prompt["22"]["inputs"]["image"] = cloth_name
|
||||
|
||||
#input human img
|
||||
prompt["61"]["inputs"]["image"] = human_name
|
||||
|
||||
#out put name
|
||||
out_img_name = str(uuid.uuid4())[:8]
|
||||
prompt["102"]["inputs"]["filename_prefix"] = out_img_name
|
||||
|
||||
p = {"prompt": prompt}
|
||||
|
||||
data = json.dumps(p).encode('utf-8')
|
||||
|
||||
response = requests.post("http://localhost:8188/prompt", data=data)
|
||||
prompt_id = response.json()["prompt_id"]
|
||||
|
||||
|
||||
# 2. 轮询队列,直到任务完成
|
||||
while True:
|
||||
queue = requests.get("http://localhost:8188/queue").json()
|
||||
# print(queue)
|
||||
if not queue["queue_running"] and not queue["queue_pending"]:
|
||||
break # 队列为空,任务已完成
|
||||
time.sleep(0.5) # 避免频繁请求
|
||||
|
||||
# 3. 从历史记录中获取结果
|
||||
history = requests.get("http://localhost:8188/history").json()
|
||||
# print("History:", history)
|
||||
outputs = history[prompt_id]["outputs"]
|
||||
for out in outputs:
|
||||
if 'images' in outputs[out]:
|
||||
for out_img in outputs[out]['images']:
|
||||
if 'filename' in out_img:
|
||||
out_img_file_name = out_img['filename']
|
||||
if out_img_name in out_img_file_name:
|
||||
return out_img_file_name, "success"
|
||||
return None, "can not find output image"
|
||||
|
||||
def save_base64_image(base64_str, prefix):
|
||||
"""
|
||||
将base64字符串保存为图片文件
|
||||
:param base64_str: 带有类型标识的base64字符串
|
||||
:param prefix: 文件名前缀
|
||||
:return: 保存的文件路径
|
||||
"""
|
||||
try:
|
||||
# 分离base64头部和实际数据
|
||||
header, data = base64_str.split(',', 1)
|
||||
|
||||
# 从头部获取文件扩展名
|
||||
file_ext = ''
|
||||
if 'image/png' in header:
|
||||
file_ext = '.png'
|
||||
elif 'image/jpeg' in header:
|
||||
file_ext = '.jpg'
|
||||
elif 'image/jpg' in header:
|
||||
file_ext = '.jpg'
|
||||
elif 'image/gif' in header:
|
||||
file_ext = '.gif'
|
||||
else:
|
||||
file_ext = '.png' # 默认使用png
|
||||
|
||||
# 生成唯一文件名
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
unique_id = str(uuid.uuid4())[:8]
|
||||
filename = f"{prefix}_{timestamp}_{unique_id}{file_ext}"
|
||||
filepath = os.path.join('/home/szlc/code/ComfyUI/change_cloth/static/imgs', filename)
|
||||
|
||||
# 解码并保存图片
|
||||
with open(filepath, 'wb') as f:
|
||||
f.write(base64.b64decode(data))
|
||||
|
||||
input_filepath = os.path.join('/home/szlc/code/ComfyUI/input', filename)
|
||||
# 解码并保存图片
|
||||
with open(input_filepath, 'wb') as f:
|
||||
f.write(base64.b64decode(data))
|
||||
|
||||
return filename
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error saving image: {str(e)}")
|
||||
return None
|
||||
|
||||
def get_image_dimensions(image_path):
|
||||
try:
|
||||
with Image.open(image_path) as img:
|
||||
width, height = img.size
|
||||
return width, height
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
return None
|
||||
|
||||
def image_to_base64(file_path, mime_type=None):
|
||||
if mime_type is None:
|
||||
extension = file_path.split('.')[-1].lower()
|
||||
mime_types = {
|
||||
'jpg': 'image/jpeg',
|
||||
'jpeg': 'image/jpeg',
|
||||
'png': 'image/png',
|
||||
'gif': 'image/gif',
|
||||
'webp': 'image/webp',
|
||||
'bmp': 'image/bmp'
|
||||
}
|
||||
mime_type = mime_types.get(extension, 'application/octet-stream')
|
||||
|
||||
# 读取文件内容并编码为Base64
|
||||
with open(file_path, 'rb') as image_file:
|
||||
encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
|
||||
|
||||
# 组合成Data URI格式
|
||||
return f"data:{mime_type};base64,{encoded_string}"
|
||||
|
||||
def upload_to_oss(image_path, object_name=None):
|
||||
# 配置信息(替换为你的实际信息)
|
||||
access_key_id = 'LTAI5tB9t2RH6f1drSLvVLLZ'
|
||||
access_key_secret = '91uzPI1RAFHN7n3Y6TJDGFP8w0dG1R'
|
||||
endpoint = 'oss-cn-beijing.aliyuncs.com' # 替换为你的Endpoint
|
||||
bucket_name = 'llyz'
|
||||
|
||||
# 创建Bucket实例
|
||||
auth = oss2.Auth(access_key_id, access_key_secret)
|
||||
bucket = oss2.Bucket(auth, endpoint, bucket_name)
|
||||
|
||||
# 如果没有指定OSS文件名,则使用本地文件名
|
||||
if object_name is None:
|
||||
object_name = image_path.split('/')[-1] # 取本地文件名
|
||||
|
||||
try:
|
||||
# 上传文件
|
||||
bucket.put_object_from_file(object_name, image_path)
|
||||
|
||||
# 获取文件URL(有效期默认10年)
|
||||
url = bucket.sign_url('GET', object_name, 3600 * 24 * 365 * 10)
|
||||
|
||||
# 或者使用公共读Bucket的URL(如果Bucket是公共读权限)
|
||||
# url = f"https://{bucket_name}.{endpoint}/{object_name}"
|
||||
|
||||
print(f"文件上传成功,URL: {url}")
|
||||
return url
|
||||
except Exception as e:
|
||||
print(f"上传失败: {str(e)}")
|
||||
return None
|
||||
|
||||
|
||||
def process_change_cloth(human_filename, cloth_filename, output_format, img_url, is_generated):
|
||||
w,h = get_image_dimensions(f'/home/szlc/code/ComfyUI/input/{human_filename}')
|
||||
|
||||
out_put_name, msg = change(human_filename, cloth_filename, w, h, img_url, is_generated)
|
||||
if out_put_name == None:
|
||||
print(f'Failed to change cloth {msg}')
|
||||
return jsonify({"ret":-1, 'msg': f'Failed to change cloth {msg}'}), 200
|
||||
|
||||
image = Image.open(f'/home/szlc/code/ComfyUI/output/{out_put_name}')
|
||||
jpg_name = out_put_name.replace(".png", ".jpg")
|
||||
jpg_path_name = f'/home/szlc/code/ComfyUI/change_cloth/static/imgs/{jpg_name}'
|
||||
image.save(jpg_path_name, quality=95)
|
||||
|
||||
|
||||
upload_to_oss(jpg_path_name, jpg_name) # 第二个参数可选,指定OSS上的路径
|
||||
https_url = f'https://llyz.oss-cn-beijing.aliyuncs.com/{jpg_name}'
|
||||
print(f"生成的HTTPS URL: {https_url}")
|
||||
|
||||
if 'base64' in output_format:
|
||||
return jsonify({
|
||||
"ret":0,
|
||||
"state": 0,
|
||||
"msg":"success",
|
||||
"data":image_to_base64(jpg_path_name)
|
||||
})
|
||||
else:
|
||||
return jsonify({
|
||||
"ret":0,
|
||||
"state": 0,
|
||||
"msg":"success",
|
||||
"url":https_url
|
||||
})
|
||||
|
||||
|
||||
def get_file_extension(url_or_filename):
|
||||
"""从 URL 或文件名中提取扩展名(如 .jpg、.png)"""
|
||||
# 处理 URL 情况(如 https://example.com/image.jpg?width=200)
|
||||
if url_or_filename.startswith(('http://', 'https://')):
|
||||
parsed = urlparse(url_or_filename)
|
||||
path = parsed.path
|
||||
else:
|
||||
path = url_or_filename
|
||||
|
||||
# 提取扩展名(转换为小写,去掉问号后的参数)
|
||||
ext = os.path.splitext(path)[1].lower().split('?')[0]
|
||||
return ext if ext else '.png' # 默认 PNG(如果无法提取)
|
||||
|
||||
def save_image_from_url(image_url):
|
||||
"""从 URL 下载图片并保留原始格式"""
|
||||
try:
|
||||
response = requests.get(image_url, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
# 获取原始图片格式
|
||||
ext = get_file_extension(image_url)
|
||||
filename = f"{uuid.uuid4()}{ext}"
|
||||
filepath = os.path.join('/home/szlc/code/ComfyUI/change_cloth/static/imgs', filename)
|
||||
|
||||
with open(filepath, "wb") as f:
|
||||
for chunk in response.iter_content(1024):
|
||||
f.write(chunk)
|
||||
|
||||
input_filepath = os.path.join('/home/szlc/code/ComfyUI/input', filename)
|
||||
shutil.copy(filepath, input_filepath)
|
||||
|
||||
return filename
|
||||
except Exception as e:
|
||||
print(f"Error saving image from URL: {e}")
|
||||
return None
|
||||
|
||||
|
||||
@app.route('/do_change_cloth', methods=['POST'])
|
||||
def do_change_cloth():
|
||||
"""从 URL 下载图片"""
|
||||
data = request.json
|
||||
|
||||
human_url = data.get("human_url")
|
||||
|
||||
human_filename = save_image_from_url(human_url)
|
||||
if not human_filename:
|
||||
return jsonify({"state":-1, "error": "Failed to download or save human image"}), 500
|
||||
|
||||
cloth_url = data.get("cloth_url")
|
||||
cloth_filename = save_image_from_url(cloth_url)
|
||||
if not cloth_filename:
|
||||
return jsonify({"state":-1, "error": "Failed to download or save image"}), 500
|
||||
|
||||
output_format = data.get('output_format')
|
||||
|
||||
return process_change_cloth(human_filename, cloth_filename, output_format, cloth_url, False)
|
||||
|
||||
get_redis_conn().set(QUEUE_NAME, "[]")
|
||||
|
||||
def pushStr2Queue(key, data_str):
|
||||
redis_conn = get_redis_conn()
|
||||
lock = acquire_lock(redis_conn, KEY_QUEUE_LOCK_NAME)
|
||||
if lock:
|
||||
try:
|
||||
redis_conn.set(key, data_str, ex=180)
|
||||
key_queue_str = redis_conn.get(QUEUE_NAME).decode("utf-8")
|
||||
key_queue = json.loads(key_queue_str)
|
||||
key_queue.append(key)
|
||||
logger.info(f"pushStr2Queue queue len:{key_queue}")
|
||||
key_queue_str = json.dumps(key_queue)
|
||||
redis_conn.set(QUEUE_NAME, key_queue_str)
|
||||
finally:
|
||||
lock.release()
|
||||
else:
|
||||
logger.error(f"pushStr2Queue get lock error {data_str}")
|
||||
|
||||
@app.route('/change_cloth', methods=['POST'])
|
||||
def change_cloth():
|
||||
"""从 URL 下载图片"""
|
||||
data = request.json
|
||||
|
||||
human_url = data.get("human_url")
|
||||
if not human_url:
|
||||
return jsonify({"state":-1, "error": "Missing 'human_url' parameter"}), 400
|
||||
|
||||
cloth_url = data.get("cloth_url")
|
||||
if not cloth_url:
|
||||
return jsonify({"state":-1, "error": "Missing 'cloth_url' parameter"}), 400
|
||||
|
||||
output_format = data.get('output_format')
|
||||
if not output_format:
|
||||
return jsonify({"state":-1, "error": "Missing 'cloth_url' parameter"}), 500
|
||||
|
||||
|
||||
key = str(uuid.uuid4())
|
||||
redis_conn = get_redis_conn()
|
||||
|
||||
task_data = {}
|
||||
task_data['key'] = key
|
||||
task_data['api'] = "/do_change_cloth"
|
||||
task_data['request'] = data
|
||||
logger.info(f"request do_change_cloth img:{data['human_url'][:64]}, output_format:{data['output_format']}")
|
||||
task_data_str = json.dumps(task_data)
|
||||
pushStr2Queue(key, task_data_str)
|
||||
timeout = DEFAULT_TIMEOUT
|
||||
start_time = datetime.now()
|
||||
result_key = f"result_{key}"
|
||||
result_status_key = f"result_status_code_{key}"
|
||||
while (datetime.now() - start_time).seconds < timeout:
|
||||
if redis_conn.exists(result_key):
|
||||
result_str = redis_conn.get(result_key)
|
||||
return result_str, int(redis_conn.get(result_status_key)), {'Content-Type': 'application/json'}
|
||||
time.sleep(0.1)
|
||||
|
||||
logger.error(f"request hairColor time out ")
|
||||
return jsonify({
|
||||
'msg': f'Timeout after {timeout} seconds, http time out'
|
||||
, "state":-1, "data":""
|
||||
}), 408
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
app.run(host="0.0.0.0", port=8888, debug=True)
|
||||
@@ -0,0 +1,39 @@
|
||||
import cv2
|
||||
import mediapipe as mp
|
||||
|
||||
def is_thigh_visible(image_path):
|
||||
# 初始化MediaPipe Pose模型
|
||||
mp_pose = mp.solutions.pose
|
||||
pose = mp_pose.Pose(static_image_mode=True, min_detection_confidence=0.5)
|
||||
|
||||
# 读取图像
|
||||
image = cv2.imread(image_path)
|
||||
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
|
||||
# 进行姿态估计
|
||||
results = pose.process(image_rgb)
|
||||
|
||||
if not results.pose_landmarks:
|
||||
return False # 未检测到人体
|
||||
|
||||
# 获取关键点
|
||||
landmarks = results.pose_landmarks.landmark
|
||||
|
||||
# 检查髋部和膝盖之间的关键点(大腿部分)
|
||||
# MediaPipe关键点索引:
|
||||
# 23: 左髋, 24: 右髋
|
||||
# 25: 左膝, 26: 右膝
|
||||
|
||||
# 检查左大腿是否可见(髋部和膝盖之间)
|
||||
left_hip_visible = landmarks[mp_pose.PoseLandmark.LEFT_HIP.value].visibility > 0.6
|
||||
left_knee_visible = landmarks[mp_pose.PoseLandmark.LEFT_KNEE.value].visibility > 0.6
|
||||
left_thigh_visible = left_hip_visible and left_knee_visible
|
||||
|
||||
# 检查右大腿是否可见
|
||||
right_hip_visible = landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value].visibility > 0.6
|
||||
right_knee_visible = landmarks[mp_pose.PoseLandmark.RIGHT_KNEE.value].visibility > 0.6
|
||||
right_thigh_visible = right_hip_visible and right_knee_visible
|
||||
|
||||
return left_thigh_visible and right_thigh_visible
|
||||
|
||||
print(is_thigh_visible("/home/szlc/code/ComfyUI/change_cloth/static/imgs/fc74eded_00001_.jpg"))
|
||||
@@ -0,0 +1,50 @@
|
||||
import time
|
||||
from redis.lock import Lock
|
||||
import logging
|
||||
|
||||
QUEUE_NAME = 'task_queue'
|
||||
# 配置
|
||||
REDIS_HOST = '127.0.0.1'
|
||||
REDIS_PORT = 6379
|
||||
REDIS_DB = 0
|
||||
DEFAULT_TIMEOUT = (60*10)
|
||||
GPU_SERVER_TIME_OUT = 180
|
||||
GPU_SERVER_LIST = 'gpu_server_list_key'
|
||||
SERVER_LOG_ = "server_log"
|
||||
SERVER_LOG_EVENT_LEN = 20
|
||||
KEY_QUEUE_LOCK_NAME = 'KEY_QUEUE_LOCK_NAME'
|
||||
SERVER_LIST_LOCK_NAME = 'SERVER_LIST_LOCK_NAME'
|
||||
SERVER_LIST_LOG_LOCK_NAME = 'SERVER_LIST_LOG_LOCK_NAME'
|
||||
|
||||
# 创建 logger
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.setLevel(logging.INFO) # 设置 logger 的级别
|
||||
|
||||
# 创建文件 handler
|
||||
file_handler = logging.FileHandler('/var/log/fuyan/config.log')
|
||||
file_handler.setLevel(logging.INFO)
|
||||
file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
|
||||
file_handler.setFormatter(file_formatter)
|
||||
|
||||
# 创建控制台 handler
|
||||
console_handler = logging.StreamHandler()
|
||||
console_handler.setLevel(logging.INFO)
|
||||
console_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
|
||||
console_handler.setFormatter(console_formatter)
|
||||
|
||||
# 添加 handlers 到 logger
|
||||
logger.addHandler(file_handler)
|
||||
logger.addHandler(console_handler)
|
||||
|
||||
def acquire_lock(redis_client, lock_name):
|
||||
lock = Lock(
|
||||
redis_client,
|
||||
name=lock_name,
|
||||
timeout=30,
|
||||
blocking_timeout=10
|
||||
)
|
||||
acquired = lock.acquire()
|
||||
if acquired:
|
||||
return lock
|
||||
return None
|
||||
|
||||
@@ -1,323 +0,0 @@
|
||||
import os
|
||||
import uuid
|
||||
import requests
|
||||
from flask import Flask, request, jsonify
|
||||
from urllib.parse import urlparse
|
||||
from volcenginesdkarkruntime import Ark
|
||||
import json
|
||||
import base64
|
||||
import time
|
||||
from PIL import Image
|
||||
from datetime import datetime
|
||||
|
||||
APP_ROOT = os.path.dirname(os.path.abspath(__file__))
|
||||
STATIC_FOLDER = os.path.join(APP_ROOT, 'static')
|
||||
|
||||
app = Flask(__name__)
|
||||
|
||||
# 请确保您已将 API Key 存储在环境变量 ARK_API_KEY 中
|
||||
# 初始化Ark客户端,从环境变量中读取您的API Key
|
||||
client = Ark(
|
||||
# 此为默认路径,您可根据业务所在地域进行配置
|
||||
base_url="https://ark.cn-beijing.volces.com/api/v3",
|
||||
# 从环境变量中获取您的 API Key。此为默认方式,您可根据需要进行修改
|
||||
# api_key=os.environ.get("ARK_API_KEY"),
|
||||
api_key='14fc0280-fc65-462d-ac2d-50178c0212e3'
|
||||
)
|
||||
|
||||
def parse_facial_data(data_str):
|
||||
print(f'原始字符串 {data_str}')
|
||||
# 去除多余的标记和换行符
|
||||
json_str_clean = data_str.strip('```json\n').strip('```').strip()
|
||||
|
||||
# 转换为Python字典
|
||||
data = json.loads(json_str_clean)
|
||||
s = data['面部年龄']
|
||||
s = s.replace('岁', '')
|
||||
parts = s.split('-')
|
||||
start_age = int(parts[0])
|
||||
end_age = int(parts[1])
|
||||
data['start_age'] = start_age
|
||||
data['end_age'] = end_age
|
||||
return data
|
||||
|
||||
def GetPicDesc(img_url):
|
||||
response = client.chat.completions.create(
|
||||
# 指定您创建的方舟推理接入点 ID,此处已帮您修改为您的推理接入点 ID
|
||||
model="doubao-1.5-vision-pro-250328",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": img_url
|
||||
},
|
||||
},
|
||||
{"type": "text", "text": "告诉我图片中的服装长度, 只要答案, 选项有(腰部以上、腰、跨、大腿上部1/3、大腿中部、膝盖、小腿、小腿上部1/3、脚踝、拖地、难以辨认)"},
|
||||
],
|
||||
}
|
||||
],
|
||||
|
||||
)
|
||||
|
||||
text = response.choices[0].message.content
|
||||
result = {}
|
||||
|
||||
try:
|
||||
result = parse_facial_data(text)
|
||||
except:
|
||||
print(f'Error {text}')
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def change(human_name, cloth_name, c_width, c_height):
|
||||
with open('/home/szlc/code/ComfyUI/change_cloth/change_new.json', 'r', encoding='utf-8') as file:
|
||||
prompt_text = file.read()
|
||||
|
||||
prompt = json.loads(prompt_text)
|
||||
|
||||
prompt["99"]["inputs"]["width"] = int((c_width/c_height) * 1024)
|
||||
|
||||
#input cloth img
|
||||
prompt["22"]["inputs"]["image"] = cloth_name
|
||||
|
||||
#input human img
|
||||
prompt["61"]["inputs"]["image"] = human_name
|
||||
|
||||
#out put name
|
||||
out_img_name = str(uuid.uuid4())[:8]
|
||||
prompt["102"]["inputs"]["filename_prefix"] = out_img_name
|
||||
|
||||
p = {"prompt": prompt}
|
||||
|
||||
data = json.dumps(p).encode('utf-8')
|
||||
|
||||
response = requests.post("http://localhost:8188/prompt", data=data)
|
||||
prompt_id = response.json()["prompt_id"]
|
||||
|
||||
|
||||
# 2. 轮询队列,直到任务完成
|
||||
while True:
|
||||
queue = requests.get("http://localhost:8188/queue").json()
|
||||
# print(queue)
|
||||
if not queue["queue_running"] and not queue["queue_pending"]:
|
||||
break # 队列为空,任务已完成
|
||||
time.sleep(0.5) # 避免频繁请求
|
||||
|
||||
# 3. 从历史记录中获取结果
|
||||
history = requests.get("http://localhost:8188/history").json()
|
||||
print("History:", history)
|
||||
outputs = history[prompt_id]["outputs"]
|
||||
for out in outputs:
|
||||
if 'images' in outputs[out]:
|
||||
for out_img_name in outputs[out]['images']:
|
||||
if 'filename' in out_img_name:
|
||||
return out_img_name['filename']
|
||||
return None
|
||||
|
||||
def save_base64_image(base64_str, prefix):
|
||||
"""
|
||||
将base64字符串保存为图片文件
|
||||
:param base64_str: 带有类型标识的base64字符串
|
||||
:param prefix: 文件名前缀
|
||||
:return: 保存的文件路径
|
||||
"""
|
||||
try:
|
||||
# 分离base64头部和实际数据
|
||||
header, data = base64_str.split(',', 1)
|
||||
|
||||
# 从头部获取文件扩展名
|
||||
file_ext = ''
|
||||
if 'image/png' in header:
|
||||
file_ext = '.png'
|
||||
elif 'image/jpeg' in header:
|
||||
file_ext = '.jpg'
|
||||
elif 'image/jpg' in header:
|
||||
file_ext = '.jpg'
|
||||
elif 'image/gif' in header:
|
||||
file_ext = '.gif'
|
||||
else:
|
||||
file_ext = '.png' # 默认使用png
|
||||
|
||||
# 生成唯一文件名
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
unique_id = str(uuid.uuid4())[:8]
|
||||
filename = f"{prefix}_{timestamp}_{unique_id}{file_ext}"
|
||||
filepath = os.path.join('/home/szlc/code/ComfyUI/change_cloth/static', filename)
|
||||
|
||||
# 解码并保存图片
|
||||
with open(filepath, 'wb') as f:
|
||||
f.write(base64.b64decode(data))
|
||||
|
||||
input_filepath = os.path.join('/home/szlc/code/ComfyUI/input', filename)
|
||||
# 解码并保存图片
|
||||
with open(input_filepath, 'wb') as f:
|
||||
f.write(base64.b64decode(data))
|
||||
|
||||
return filename
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error saving image: {str(e)}")
|
||||
return None
|
||||
|
||||
def get_image_dimensions(image_path):
|
||||
try:
|
||||
with Image.open(image_path) as img:
|
||||
width, height = img.size
|
||||
return width, height
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
return None
|
||||
|
||||
def image_to_base64(file_path, mime_type=None):
|
||||
if mime_type is None:
|
||||
extension = file_path.split('.')[-1].lower()
|
||||
mime_types = {
|
||||
'jpg': 'image/jpeg',
|
||||
'jpeg': 'image/jpeg',
|
||||
'png': 'image/png',
|
||||
'gif': 'image/gif',
|
||||
'webp': 'image/webp',
|
||||
'bmp': 'image/bmp'
|
||||
}
|
||||
mime_type = mime_types.get(extension, 'application/octet-stream')
|
||||
|
||||
# 读取文件内容并编码为Base64
|
||||
with open(file_path, 'rb') as image_file:
|
||||
encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
|
||||
|
||||
# 组合成Data URI格式
|
||||
return f"data:{mime_type};base64,{encoded_string}"
|
||||
|
||||
def process_change_cloth(human_filename, cloth_filename, output_format):
|
||||
w,h = get_image_dimensions(f'/home/szlc/code/ComfyUI/input/{human_filename}')
|
||||
|
||||
out_put_name = change(human_filename, cloth_filename, w, h)
|
||||
if out_put_name == None:
|
||||
return jsonify({"ret":-1, 'msg': 'Failed to change cloth'}), 500
|
||||
|
||||
image = Image.open(f'/home/szlc/code/ComfyUI/output/{out_put_name}')
|
||||
jpg_name = out_put_name.replace(".png", ".jpg")
|
||||
jpg_path_name = f'/home/szlc/code/ComfyUI/change_cloth/static/{jpg_name}'
|
||||
image.save(jpg_path_name, quality=95)
|
||||
|
||||
if 'base64' in output_format:
|
||||
return jsonify({
|
||||
"ret":0,
|
||||
"msg":"success",
|
||||
"data":image_to_base64(jpg_path_name)
|
||||
})
|
||||
else:
|
||||
return jsonify({
|
||||
"ret":0,
|
||||
"msg":"success",
|
||||
"url":f"http://112.126.94.241:18018/static/{jpg_name}"
|
||||
})
|
||||
|
||||
@app.route('/change_cloth_base64', methods=['POST'])
|
||||
def change_cloth_base64():
|
||||
# 获取参数
|
||||
data = request.get_json()
|
||||
if not data:
|
||||
return jsonify({"ret":-1,'msg': 'No JSON data provided'}), 400
|
||||
|
||||
human_img = data.get('human_img')
|
||||
cloth_img = data.get('cloth_img')
|
||||
output_format = data.get('output_format')
|
||||
|
||||
if not human_img or not cloth_img:
|
||||
return jsonify({"ret":-1, 'msg': 'Both human_img and cloth_img are required'}), 400
|
||||
|
||||
try:
|
||||
|
||||
if 'base64' in output_format:
|
||||
return jsonify({
|
||||
"ret":0,
|
||||
"msg":"success",
|
||||
"data":image_to_base64("/home/szlc/code/ComfyUI/input/girl_full.jpg")
|
||||
})
|
||||
else:
|
||||
return jsonify({
|
||||
"ret":0,
|
||||
"msg":"success",
|
||||
"url":f"http://43.143.205.217/imgs/processed_f60d0a9862834e1c92bc1c64ecf8f3ba.jpg"
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
return jsonify({"ret":-1, 'error': str(e)}), 500
|
||||
|
||||
def get_file_extension(url_or_filename):
|
||||
"""从 URL 或文件名中提取扩展名(如 .jpg、.png)"""
|
||||
# 处理 URL 情况(如 https://example.com/image.jpg?width=200)
|
||||
if url_or_filename.startswith(('http://', 'https://')):
|
||||
parsed = urlparse(url_or_filename)
|
||||
path = parsed.path
|
||||
else:
|
||||
path = url_or_filename
|
||||
|
||||
# 提取扩展名(转换为小写,去掉问号后的参数)
|
||||
ext = os.path.splitext(path)[1].lower().split('?')[0]
|
||||
return ext if ext else '.png' # 默认 PNG(如果无法提取)
|
||||
|
||||
def save_image_from_url(image_url):
|
||||
"""从 URL 下载图片并保留原始格式"""
|
||||
try:
|
||||
response = requests.get(image_url, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
# 获取原始图片格式
|
||||
ext = get_file_extension(image_url)
|
||||
filename = f"{uuid.uuid4()}{ext}"
|
||||
filepath = os.path.join('/home/szlc/code/ComfyUI/change_cloth/static', filename)
|
||||
|
||||
with open(filepath, "wb") as f:
|
||||
for chunk in response.iter_content(1024):
|
||||
f.write(chunk)
|
||||
|
||||
input_filepath = os.path.join('/home/szlc/code/ComfyUI/input', filename)
|
||||
with open(input_filepath, "wb") as f:
|
||||
for chunk in response.iter_content(1024):
|
||||
f.write(chunk)
|
||||
|
||||
return filename
|
||||
except Exception as e:
|
||||
print(f"Error saving image from URL: {e}")
|
||||
return None
|
||||
|
||||
@app.route('/change_cloth', methods=['POST'])
|
||||
def change_cloth():
|
||||
"""从 URL 下载图片"""
|
||||
data = request.json
|
||||
human_url = data.get("human_url")
|
||||
|
||||
if not human_url:
|
||||
return jsonify({"error": "Missing 'human_url' parameter"}), 400
|
||||
|
||||
|
||||
cloth_url = data.get("cloth_url")
|
||||
|
||||
if not cloth_url:
|
||||
return jsonify({"error": "Missing 'cloth_url' parameter"}), 400
|
||||
|
||||
output_format = data.get('output_format')
|
||||
if not output_format:
|
||||
return jsonify({"error": "Missing 'output_format' parameter"}), 500
|
||||
|
||||
if 'base64' in output_format:
|
||||
return jsonify({
|
||||
"ret":0,
|
||||
"msg":"success",
|
||||
"data":image_to_base64("/home/szlc/code/ComfyUI/input/girl_full.jpg")
|
||||
})
|
||||
else:
|
||||
return jsonify({
|
||||
"ret":0,
|
||||
"msg":"success",
|
||||
"url":f"http://43.143.205.217/imgs/processed_f60d0a9862834e1c92bc1c64ecf8f3ba.jpg"
|
||||
})
|
||||
|
||||
if __name__ == '__main__':
|
||||
app.run(host="0.0.0.0", port=8888, debug=False)
|
||||
+84
@@ -0,0 +1,84 @@
|
||||
import time
|
||||
import logging
|
||||
import redis
|
||||
import uuid
|
||||
import json
|
||||
from flask import Flask, request, jsonify, send_from_directory
|
||||
from datetime import datetime, timedelta
|
||||
from config import QUEUE_NAME, DEFAULT_TIMEOUT
|
||||
from flask_cors import CORS
|
||||
GPU_SERVER_LIST = 'gpu_server_list_key'
|
||||
|
||||
# 创建 logger
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.setLevel(logging.INFO) # 设置 logger 的级别
|
||||
|
||||
# 创建文件 handler
|
||||
file_handler = logging.FileHandler('/var/log/fuyan/monitor.log')
|
||||
file_handler.setLevel(logging.INFO)
|
||||
file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
|
||||
file_handler.setFormatter(file_formatter)
|
||||
|
||||
# 创建控制台 handler
|
||||
console_handler = logging.StreamHandler()
|
||||
console_handler.setLevel(logging.INFO)
|
||||
console_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
|
||||
console_handler.setFormatter(console_formatter)
|
||||
|
||||
# 添加 handlers 到 logger
|
||||
logger.addHandler(file_handler)
|
||||
logger.addHandler(console_handler)
|
||||
|
||||
monitor = Flask(__name__)
|
||||
CORS(monitor)
|
||||
|
||||
# 配置
|
||||
from config import REDIS_HOST
|
||||
from config import REDIS_PORT
|
||||
from config import REDIS_DB
|
||||
from config import SERVER_LOG_
|
||||
|
||||
# 创建 Redis 连接池
|
||||
redis_pool = redis.ConnectionPool(
|
||||
host=REDIS_HOST,
|
||||
port=REDIS_PORT,
|
||||
db=REDIS_DB,
|
||||
max_connections=2000 # 根据实际情况调整
|
||||
)
|
||||
|
||||
def get_redis_conn():
|
||||
"""获取 Redis 连接"""
|
||||
return redis.Redis(connection_pool=redis_pool)
|
||||
|
||||
@monitor.route('/get_server_state', methods=['GET', 'POST'])
|
||||
def get_server_state():
|
||||
monitor_data = []
|
||||
redis_conn = get_redis_conn()
|
||||
server_list_str = redis_conn.get(GPU_SERVER_LIST).decode("utf-8")
|
||||
server_list = json.loads(server_list_str)
|
||||
for name in server_list:
|
||||
server_str = redis_conn.get(name).decode("utf-8")
|
||||
server = json.loads(server_str)
|
||||
|
||||
server_log_name = f"{SERVER_LOG_}{name}"
|
||||
server_log_str = redis_conn.get(server_log_name).decode("utf-8")
|
||||
server_log = json.loads(server_log_str)
|
||||
server_log['busy'] = not server['can_use']
|
||||
monitor_data.append(server_log)
|
||||
json_str = json.dumps(monitor_data)
|
||||
print(len(json_str))
|
||||
return json_str, 200
|
||||
|
||||
|
||||
@monitor.route('/')
|
||||
def serve_index():
|
||||
return send_from_directory('static', 'index.html')
|
||||
|
||||
if __name__ == '__main__':
|
||||
# 启动Flask应用,启用多线程处理
|
||||
monitor.run(
|
||||
host='0.0.0.0',
|
||||
port=5001,
|
||||
threaded=True, # 启用多线程处理并发请求
|
||||
debug=False # 生产环境应设置为False
|
||||
)
|
||||
@@ -0,0 +1,321 @@
|
||||
import time
|
||||
import datetime
|
||||
import uuid
|
||||
import redis
|
||||
import logging
|
||||
import requests
|
||||
import threading
|
||||
import json
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
from config import REDIS_HOST
|
||||
from config import REDIS_PORT
|
||||
from config import REDIS_DB
|
||||
from config import QUEUE_NAME
|
||||
from config import GPU_SERVER_LIST
|
||||
from config import SERVER_LOG_EVENT_LEN
|
||||
from config import SERVER_LOG_
|
||||
from config import GPU_SERVER_TIME_OUT
|
||||
from config import KEY_QUEUE_LOCK_NAME, acquire_lock, SERVER_LIST_LOCK_NAME, SERVER_LIST_LOG_LOCK_NAME
|
||||
import copy
|
||||
|
||||
import logging
|
||||
import csv
|
||||
|
||||
# 创建 logger
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.setLevel(logging.INFO) # 设置 logger 的级别
|
||||
|
||||
|
||||
# 获取当前日期和时间
|
||||
now = datetime.now()
|
||||
|
||||
# 格式化为字符串(例如:2023-10-25 14:30:45)
|
||||
date_time_str = now.strftime("%Y-%m-%d_%H:%M:%S")
|
||||
|
||||
# 创建文件 handler
|
||||
file_handler = logging.FileHandler(f'/var/log/fuyan/worker_{date_time_str}.log')
|
||||
file_handler.setLevel(logging.INFO)
|
||||
file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
|
||||
file_handler.setFormatter(file_formatter)
|
||||
|
||||
# 创建控制台 handler
|
||||
console_handler = logging.StreamHandler()
|
||||
console_handler.setLevel(logging.INFO)
|
||||
console_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
|
||||
console_handler.setFormatter(console_formatter)
|
||||
|
||||
# 添加 handlers 到 logger
|
||||
logger.addHandler(file_handler)
|
||||
logger.addHandler(console_handler)
|
||||
|
||||
# 创建 Redis 连接池
|
||||
redis_pool = redis.ConnectionPool(
|
||||
host=REDIS_HOST,
|
||||
port=REDIS_PORT,
|
||||
db=REDIS_DB,
|
||||
max_connections=2000 # 根据实际情况调整
|
||||
)
|
||||
|
||||
def get_redis_conn():
|
||||
"""获取 Redis 连接"""
|
||||
return redis.Redis(connection_pool=redis_pool)
|
||||
|
||||
redis_conn = get_redis_conn()
|
||||
|
||||
def get_time(timestamp):
|
||||
# 转换为UTC时间
|
||||
utc_time = datetime.utcfromtimestamp(timestamp)
|
||||
|
||||
# 转换为北京时间 (UTC+8)
|
||||
beijing_time = utc_time + timedelta(hours=8) # 正确
|
||||
|
||||
# 格式化输出
|
||||
formatted_time = beijing_time.strftime("%Y-%m-%d %H:%M:%S")
|
||||
return formatted_time
|
||||
|
||||
def server_log_event(name, action, in_data, success=True):
|
||||
|
||||
log_data = copy.deepcopy(in_data)
|
||||
server_log_name = f"{SERVER_LOG_}{name}"
|
||||
|
||||
lock = acquire_lock(redis_conn, f"server_log_lock_{name}")
|
||||
if lock:
|
||||
try:
|
||||
logger.info(f"server_log_event name:{name} action:{action} success:{success}")
|
||||
server_log_str = redis_conn.get(server_log_name).decode("utf-8")
|
||||
server_log = json.loads(server_log_str)
|
||||
event = {}
|
||||
t = time.time()
|
||||
event['time'] = t
|
||||
event['time_str'] = get_time(t)
|
||||
event['action'] = action
|
||||
|
||||
server_log['last_call'] = "success"
|
||||
if 'state' in log_data:
|
||||
if log_data['state'] != 0:
|
||||
server_log['last_call'] = "failure"
|
||||
|
||||
if not success:
|
||||
server_log['last_call'] = "failure"
|
||||
|
||||
event['data'] = log_data
|
||||
if(len(server_log['events']) > SERVER_LOG_EVENT_LEN):
|
||||
server_log['events'].pop(0)
|
||||
server_log['events'].append(event)
|
||||
server_log['last_event'] = event
|
||||
server_log_str = json.dumps(server_log)
|
||||
redis_conn.set(f"{SERVER_LOG_}{name}", server_log_str)
|
||||
|
||||
finally:
|
||||
lock.release()
|
||||
else:
|
||||
logger.error(f"get lock none {server_log_name}")
|
||||
|
||||
def registerGpuServer(name, url, can_use):
|
||||
try:
|
||||
server_list_str = redis_conn.get(GPU_SERVER_LIST).decode("utf-8")
|
||||
server_list = json.loads(server_list_str)
|
||||
if name not in server_list:
|
||||
server_list.append(name)
|
||||
server = {}
|
||||
server['name'] = name
|
||||
server['url'] = url
|
||||
server['can_use'] = can_use
|
||||
server_str = json.dumps(server)
|
||||
logger.info(f"registerGpuServer, {server_str}")
|
||||
redis_conn.set(name, server_str)
|
||||
|
||||
server_log = {}
|
||||
server_log['name'] = name
|
||||
event = {}
|
||||
t = time.time()
|
||||
event['time'] = t
|
||||
event['time_str'] = get_time(t)
|
||||
event['action'] = "register"
|
||||
event['data'] = f"url {url}"
|
||||
server_log['last_event'] = event
|
||||
server_log['events'] = []
|
||||
server_log['last_call'] = "success"
|
||||
if(len(server_log['events']) > SERVER_LOG_EVENT_LEN):
|
||||
server_log['events'].pop(0)
|
||||
server_log['events'].append(event)
|
||||
server_log_str = json.dumps(server_log)
|
||||
redis_conn.set(f"{SERVER_LOG_}{name}", server_log_str)
|
||||
|
||||
redis_conn.set(GPU_SERVER_LIST, json.dumps(server_list))
|
||||
except Exception as e:
|
||||
logger.error(f"registerGpuServer{e.message}")
|
||||
|
||||
|
||||
def get_remote_gpu_server():
|
||||
start_time = datetime.now()
|
||||
while (datetime.now() - start_time).seconds < GPU_SERVER_TIME_OUT:
|
||||
lock = acquire_lock(redis_conn, SERVER_LIST_LOCK_NAME)
|
||||
try:
|
||||
server_list_str = redis_conn.get(GPU_SERVER_LIST).decode("utf-8")
|
||||
server_list = json.loads(server_list_str)
|
||||
for name in server_list:
|
||||
server_str = redis_conn.get(name).decode("utf-8")
|
||||
server = json.loads(server_str)
|
||||
if server['can_use']:
|
||||
return server
|
||||
finally:
|
||||
lock.release()
|
||||
time.sleep(0.2)
|
||||
return None
|
||||
|
||||
def call_remote_gpu_server(task_data_str, server=None):
|
||||
task_data = json.loads(task_data_str)
|
||||
key = task_data['key']
|
||||
if server == None:
|
||||
server = get_remote_gpu_server()
|
||||
|
||||
result_str = ''
|
||||
result_status_code = '200'
|
||||
if server:
|
||||
headers = {
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
|
||||
if 'do_change_cloth' in task_data['api']:
|
||||
data = {
|
||||
"human_url": task_data['request']['human_url'],
|
||||
"cloth_url": task_data['request']["cloth_url"],
|
||||
"output_format":task_data['request']['output_format']
|
||||
}
|
||||
|
||||
name = server['name']
|
||||
|
||||
# set server busy
|
||||
server['can_use'] = False
|
||||
server_str = json.dumps(server)
|
||||
redis_conn.set(name, server_str)
|
||||
|
||||
server_log_event(server['name'], "Request_GpuServer", data, True)
|
||||
url = server['url'] + task_data['api']
|
||||
try:
|
||||
logger.info(f"call get_remote_gpu_server {url}")
|
||||
start_ms = int(time.time() * 1000) # 转换为毫秒级整数
|
||||
base64 = False
|
||||
if data['output_format'] == 'base64':
|
||||
base64 = True
|
||||
|
||||
response = requests.post(url, headers=headers, json=data, timeout=30)
|
||||
|
||||
logger.info(f"call finshed {response.status_code} {response.headers} {response.text[:128]}")
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
end_ms = int(time.time() * 1000)
|
||||
if 'do_change_cloth' in task_data['api']:
|
||||
if base64:
|
||||
result['result'] = f"data:image/jpeg;base64,{result['result']}"
|
||||
|
||||
result['process_time_ms'] = (end_ms - start_ms)
|
||||
server_log_event(server['name'], "GpuServer_Response", result, True)
|
||||
result_str = json.dumps(result)
|
||||
result_status_code = '200'
|
||||
if 'state' in result:
|
||||
if result['state'] != 0:
|
||||
result_status_code = '500'
|
||||
else:
|
||||
try:
|
||||
result = response.json()
|
||||
except json.JSONDecodeError:
|
||||
result = {
|
||||
"state":-1,
|
||||
"ret":-1,
|
||||
"data":"",
|
||||
'msg': f'status_code error{response.status_code} call_remote_gpu_server{url} {response.text}'
|
||||
}
|
||||
|
||||
result_str = json.dumps(result)
|
||||
result_status_code = str(response.status_code)
|
||||
server_log_event(server['name'], f"error GpuServer_Response ", result, False)
|
||||
except Exception as e:
|
||||
logger.info(f"End call Exception {url} {type(e).__name__}, 错误信息: {e}")
|
||||
result = {
|
||||
"state":-1,
|
||||
"ret":-1,
|
||||
"data":"",
|
||||
'msg': f'Exception call call_remote_gpu_server{url}'
|
||||
}
|
||||
result_str = json.dumps(result)
|
||||
result_status_code = '500'
|
||||
server_log_event(server['name'], f"error GpuServer_Response ", result, False)
|
||||
|
||||
logger.info(f"End call {url}")
|
||||
# release server
|
||||
server['can_use'] = True
|
||||
server_str = json.dumps(server)
|
||||
redis_conn.set(name, server_str)
|
||||
|
||||
else:
|
||||
result_str = json.dumps({
|
||||
"state":-1,
|
||||
"ret":-1,
|
||||
"data":"",
|
||||
'msg': f'Timeout after all gpu server is busy'
|
||||
})
|
||||
result_status_code = '500'
|
||||
|
||||
result_key = f"result_{key}"
|
||||
redis_conn.set(result_key, result_str, ex=60)
|
||||
result_status_key = f"result_status_code_{key}"
|
||||
redis_conn.set(result_status_key, result_status_code, ex=60)
|
||||
|
||||
|
||||
def get_task_queue_key():
|
||||
redis_conn = get_redis_conn()
|
||||
lock = acquire_lock(redis_conn, KEY_QUEUE_LOCK_NAME)
|
||||
if not lock:
|
||||
logger.error("get_task_queue_key get lock error")
|
||||
return
|
||||
try:
|
||||
key_queue_str = redis_conn.get(QUEUE_NAME).decode("utf-8")
|
||||
key_queue = json.loads(key_queue_str)
|
||||
queue_len = len(key_queue)
|
||||
if(queue_len > 0):
|
||||
logger.info(f"get_task_queue_key queue len:{queue_len}")
|
||||
key = key_queue.pop(0)
|
||||
key_queue_str = json.dumps(key_queue)
|
||||
logger.info(f"get_task_queue_key queue len:{len(key_queue)}")
|
||||
redis_conn.set(QUEUE_NAME, key_queue_str)
|
||||
return key
|
||||
else:
|
||||
return None
|
||||
finally:
|
||||
lock.release()
|
||||
|
||||
def main_worker():
|
||||
while True:
|
||||
try:
|
||||
key = get_task_queue_key()
|
||||
if key:
|
||||
task_data_str = redis_conn.get(key)
|
||||
if task_data_str:
|
||||
threading.Thread(target=call_remote_gpu_server, args=(task_data_str,)).start()
|
||||
else:
|
||||
logging.error(f"main_worker get task_data_str error with key{key}")
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.info(f"main_worker {e}")
|
||||
# check_server_work()
|
||||
time.sleep(0.1)
|
||||
|
||||
serverindex = 0
|
||||
def regServer(instance):
|
||||
global serverindex
|
||||
registerGpuServer(f"s{serverindex}_1_{instance}", f"http://localhost:8888", True)
|
||||
serverindex += 1
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
import os
|
||||
print(f"[Worker] Starting with PID: {os.getpid()}")
|
||||
redis_conn.set(GPU_SERVER_LIST, "[]")
|
||||
|
||||
servers = regServer("localhost")
|
||||
for s in servers:
|
||||
regServer(s)
|
||||
main_worker()
|
||||
@@ -0,0 +1,345 @@
|
||||
{
|
||||
"6": {
|
||||
"inputs": {
|
||||
"text": "full body shot(1.9),complete figure,front-facing,model pose, Hands are at the sides of the body, smile,clear skin glow,wearing black micro skirt,simple white background,clean composition,32k,high detail,professional beauty photography style,\n",
|
||||
"clip": [
|
||||
"68",
|
||||
1
|
||||
]
|
||||
},
|
||||
"class_type": "CLIPTextEncode",
|
||||
"_meta": {
|
||||
"title": "CLIP Text Encode (Prompt)"
|
||||
}
|
||||
},
|
||||
"10": {
|
||||
"inputs": {
|
||||
"vae_name": "ae.safetensors"
|
||||
},
|
||||
"class_type": "VAELoader",
|
||||
"_meta": {
|
||||
"title": "Load VAE"
|
||||
}
|
||||
},
|
||||
"16": {
|
||||
"inputs": {
|
||||
"sampler_name": "heun"
|
||||
},
|
||||
"class_type": "KSamplerSelect",
|
||||
"_meta": {
|
||||
"title": "KSamplerSelect"
|
||||
}
|
||||
},
|
||||
"17": {
|
||||
"inputs": {
|
||||
"scheduler": "beta",
|
||||
"steps": 20,
|
||||
"denoise": 1,
|
||||
"model": [
|
||||
"63",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "BasicScheduler",
|
||||
"_meta": {
|
||||
"title": "BasicScheduler"
|
||||
}
|
||||
},
|
||||
"25": {
|
||||
"inputs": {
|
||||
"noise_seed": 873089317273353
|
||||
},
|
||||
"class_type": "RandomNoise",
|
||||
"_meta": {
|
||||
"title": "RandomNoise"
|
||||
}
|
||||
},
|
||||
"26": {
|
||||
"inputs": {
|
||||
"guidance": 10,
|
||||
"conditioning": [
|
||||
"6",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "FluxGuidance",
|
||||
"_meta": {
|
||||
"title": "FluxGuidance"
|
||||
}
|
||||
},
|
||||
"27": {
|
||||
"inputs": {
|
||||
"width": 960,
|
||||
"height": 1280,
|
||||
"batch_size": 1
|
||||
},
|
||||
"class_type": "EmptySD3LatentImage",
|
||||
"_meta": {
|
||||
"title": "EmptySD3LatentImage"
|
||||
}
|
||||
},
|
||||
"45": {
|
||||
"inputs": {
|
||||
"pulid_file": "pulid_flux_v0.9.1.safetensors"
|
||||
},
|
||||
"class_type": "PulidFluxModelLoader",
|
||||
"_meta": {
|
||||
"title": "Load PuLID Flux Model"
|
||||
}
|
||||
},
|
||||
"47": {
|
||||
"inputs": {
|
||||
"model": [
|
||||
"62",
|
||||
0
|
||||
],
|
||||
"conditioning": [
|
||||
"26",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "BasicGuider",
|
||||
"_meta": {
|
||||
"title": "BasicGuider"
|
||||
}
|
||||
},
|
||||
"48": {
|
||||
"inputs": {
|
||||
"noise": [
|
||||
"25",
|
||||
0
|
||||
],
|
||||
"guider": [
|
||||
"47",
|
||||
0
|
||||
],
|
||||
"sampler": [
|
||||
"16",
|
||||
0
|
||||
],
|
||||
"sigmas": [
|
||||
"17",
|
||||
0
|
||||
],
|
||||
"latent_image": [
|
||||
"27",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "SamplerCustomAdvanced",
|
||||
"_meta": {
|
||||
"title": "SamplerCustomAdvanced"
|
||||
}
|
||||
},
|
||||
"49": {
|
||||
"inputs": {
|
||||
"samples": [
|
||||
"48",
|
||||
0
|
||||
],
|
||||
"vae": [
|
||||
"10",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "VAEDecode",
|
||||
"_meta": {
|
||||
"title": "VAE Decode"
|
||||
}
|
||||
},
|
||||
"51": {
|
||||
"inputs": {},
|
||||
"class_type": "PulidFluxEvaClipLoader",
|
||||
"_meta": {
|
||||
"title": "Load Eva Clip (PuLID Flux)"
|
||||
}
|
||||
},
|
||||
"53": {
|
||||
"inputs": {
|
||||
"provider": "CPU"
|
||||
},
|
||||
"class_type": "PulidFluxInsightFaceLoader",
|
||||
"_meta": {
|
||||
"title": "Load InsightFace (PuLID Flux)"
|
||||
}
|
||||
},
|
||||
"62": {
|
||||
"inputs": {
|
||||
"weight": 1,
|
||||
"start_at": 0,
|
||||
"end_at": 1,
|
||||
"model": [
|
||||
"68",
|
||||
0
|
||||
],
|
||||
"pulid_flux": [
|
||||
"45",
|
||||
0
|
||||
],
|
||||
"eva_clip": [
|
||||
"51",
|
||||
0
|
||||
],
|
||||
"face_analysis": [
|
||||
"53",
|
||||
0
|
||||
],
|
||||
"image": [
|
||||
"93",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "ApplyPulidFlux",
|
||||
"_meta": {
|
||||
"title": "Apply PuLID Flux"
|
||||
}
|
||||
},
|
||||
"63": {
|
||||
"inputs": {
|
||||
"unet_name": "flux1-dev-fp8.safetensors",
|
||||
"weight_dtype": "fp8_e4m3fn"
|
||||
},
|
||||
"class_type": "UNETLoader",
|
||||
"_meta": {
|
||||
"title": "Load Diffusion Model"
|
||||
}
|
||||
},
|
||||
"64": {
|
||||
"inputs": {
|
||||
"clip_name1": "t5xxl_fp8_e4m3fn.safetensors",
|
||||
"clip_name2": "ViT-L-14-TEXT-detail-improved-hiT-GmP-HF.safetensors",
|
||||
"type": "flux",
|
||||
"device": "default"
|
||||
},
|
||||
"class_type": "DualCLIPLoader",
|
||||
"_meta": {
|
||||
"title": "DualCLIPLoader"
|
||||
}
|
||||
},
|
||||
"65": {
|
||||
"inputs": {
|
||||
"filename_prefix": "ComfyUI",
|
||||
"images": [
|
||||
"49",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "SaveImage",
|
||||
"_meta": {
|
||||
"title": "Save Image"
|
||||
}
|
||||
},
|
||||
"68": {
|
||||
"inputs": {
|
||||
"PowerLoraLoaderHeaderWidget": {
|
||||
"type": "PowerLoraLoaderHeaderWidget"
|
||||
},
|
||||
"➕ Add Lora": "",
|
||||
"model": [
|
||||
"63",
|
||||
0
|
||||
],
|
||||
"clip": [
|
||||
"64",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "Power Lora Loader (rgthree)",
|
||||
"_meta": {
|
||||
"title": "Power Lora Loader (rgthree)"
|
||||
}
|
||||
},
|
||||
"86": {
|
||||
"inputs": {
|
||||
"anything": [
|
||||
"17",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "easy cleanGpuUsed",
|
||||
"_meta": {
|
||||
"title": "Clean VRAM Used"
|
||||
}
|
||||
},
|
||||
"87": {
|
||||
"inputs": {
|
||||
"anything": [
|
||||
"68",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "easy cleanGpuUsed",
|
||||
"_meta": {
|
||||
"title": "Clean VRAM Used"
|
||||
}
|
||||
},
|
||||
"88": {
|
||||
"inputs": {
|
||||
"anything": [
|
||||
"62",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "easy cleanGpuUsed",
|
||||
"_meta": {
|
||||
"title": "Clean VRAM Used"
|
||||
}
|
||||
},
|
||||
"89": {
|
||||
"inputs": {
|
||||
"anything": [
|
||||
"47",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "easy cleanGpuUsed",
|
||||
"_meta": {
|
||||
"title": "Clean VRAM Used"
|
||||
}
|
||||
},
|
||||
"90": {
|
||||
"inputs": {
|
||||
"anything": [
|
||||
"48",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "easy cleanGpuUsed",
|
||||
"_meta": {
|
||||
"title": "Clean VRAM Used"
|
||||
}
|
||||
},
|
||||
"91": {
|
||||
"inputs": {
|
||||
"anything": [
|
||||
"49",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "easy cleanGpuUsed",
|
||||
"_meta": {
|
||||
"title": "Clean VRAM Used"
|
||||
}
|
||||
},
|
||||
"92": {
|
||||
"inputs": {
|
||||
"anything": [
|
||||
"49",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "easy clearCacheAll",
|
||||
"_meta": {
|
||||
"title": "Clear Cache All"
|
||||
}
|
||||
},
|
||||
"93": {
|
||||
"inputs": {
|
||||
"image": "v2-2ab29473f9598208eb189be01f122a8b_r.jpg"
|
||||
},
|
||||
"class_type": "LoadImage",
|
||||
"_meta": {
|
||||
"title": "Load Image"
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user