save code

This commit is contained in:
colomi
2025-07-20 21:43:34 +08:00
parent ae71e3d540
commit b30efd013b
7 changed files with 1395 additions and 323 deletions
+556
View File
@@ -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 bodysmile,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)
+39
View File
@@ -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"))
+50
View File
@@ -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
-323
View File
@@ -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
View File
@@ -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
)
+321
View File
@@ -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()
+345
View File
@@ -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"
}
}
}