save code

This commit is contained in:
colomi
2025-07-13 09:55:49 +08:00
parent 92b1c9b3c6
commit bcebd31667
4 changed files with 582 additions and 304 deletions
+61 -71
View File
@@ -3,83 +3,70 @@ 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
from volcenginesdkarkruntime import Ark
import shutil
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
# 此为默认路径,您可根据业务所在地域进行配置
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
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、脚踝、拖地、难以辨认)"},
],
}
],
)
{"type": "text", "text": "图片是一件服装的照片,告诉我图片中的服装穿到正常人身上的长度会覆盖到人体的哪个部位, 只要答案, 选项有(胸、腰、跨、大腿、膝盖、小腿、脚踝、拖地、难以辨认)"},
],
}
],
)
text = response.choices[0].message.content
result = {}
text = response.choices[0].message.content
print(text)
return text
try:
result = parse_facial_data(text)
except:
print(f'Error {text}')
def change(human_name, cloth_name, c_width, c_height, cloth_url):
return result
cloth_len = GetPicDesc(cloth_url)
if cloth_len == '难以辨认':
return None, "get image type error"
queue = requests.get("http://localhost:8188/queue").json()
if queue["queue_running"] or queue["queue_pending"]:
return None, "cur gpu is busy"
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)
prompt["96"]["inputs"]["cloth_len"] = cloth_len
# prompt["99"]["inputs"]["width"] = int((c_width/c_height) * 1024)
#input cloth img
prompt["22"]["inputs"]["image"] = cloth_name
@@ -109,14 +96,16 @@ def change(human_name, cloth_name, c_width, c_height):
# 3. 从历史记录中获取结果
history = requests.get("http://localhost:8188/history").json()
print("History:", history)
# 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
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):
"""
@@ -146,7 +135,7 @@ def save_base64_image(base64_str, prefix):
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)
filepath = os.path.join('/home/szlc/code/ComfyUI/change_cloth/static/imgs', filename)
# 解码并保存图片
with open(filepath, 'wb') as f:
@@ -192,16 +181,17 @@ def image_to_base64(file_path, mime_type=None):
# 组合成Data URI格式
return f"data:{mime_type};base64,{encoded_string}"
def process_change_cloth(human_filename, cloth_filename, output_format):
def process_change_cloth(human_filename, cloth_filename, output_format, img_url):
w,h = get_image_dimensions(f'/home/szlc/code/ComfyUI/input/{human_filename}')
out_put_name = change(human_filename, cloth_filename, w, h)
out_put_name, msg = change(human_filename, cloth_filename, w, h, img_url)
if out_put_name == None:
return jsonify({"ret":-1, 'msg': 'Failed to change cloth'}), 500
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/{jpg_name}'
jpg_path_name = f'/home/szlc/code/ComfyUI/change_cloth/static/imgs/{jpg_name}'
image.save(jpg_path_name, quality=95)
if 'base64' in output_format:
@@ -214,7 +204,7 @@ def process_change_cloth(human_filename, cloth_filename, output_format):
return jsonify({
"ret":0,
"msg":"success",
"url":f"http://112.126.94.241:18018/static/{jpg_name}"
"url":f"http://112.126.94.241:18888/static/imgs/{jpg_name}"
})
@app.route('/change_cloth_base64', methods=['POST'])
@@ -242,7 +232,9 @@ def change_cloth_base64():
if not cloth_filename:
return jsonify({"ret":-1, 'msg': 'Failed to save cloth image'}), 500
return process_change_cloth(human_filename, cloth_filename, output_format)
img_url = f"http://112.126.94.241:18888/static/imgs/{cloth_filename}"
return process_change_cloth(human_filename, cloth_filename, output_format, img_url)
except Exception as e:
return jsonify({"ret":-1, 'error': str(e)}), 500
@@ -269,16 +261,14 @@ def save_image_from_url(image_url):
# 获取原始图片格式
ext = get_file_extension(image_url)
filename = f"{uuid.uuid4()}{ext}"
filepath = os.path.join('/home/szlc/code/ComfyUI/change_cloth/static', filename)
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)
with open(input_filepath, "wb") as f:
for chunk in response.iter_content(1024):
f.write(chunk)
shutil.copy(filepath, input_filepath)
return filename
except Exception as e:
@@ -309,7 +299,7 @@ def change_cloth():
output_format = data.get('output_format')
return process_change_cloth(human_filename, cloth_filename, output_format)
return process_change_cloth(human_filename, cloth_filename, output_format, cloth_url)
if __name__ == '__main__':
app.run(host="0.0.0.0", port=8018, debug=True)
app.run(host="0.0.0.0", port=8888, debug=True)