174 lines
7.2 KiB
Python
174 lines
7.2 KiB
Python
import json
|
||
import os
|
||
import time
|
||
import base64
|
||
|
||
from JiMeng.Kit.JmCommon import JmCommon
|
||
from JiMeng.Kit.JmErrorCode import JmErrorCode
|
||
|
||
|
||
class JmTxt2Img:
|
||
action = "CVProcess"
|
||
# V2.1
|
||
req_key = "jimeng_high_aes_general_v21_L"
|
||
# V3.0
|
||
# req_key="jimeng_t2i_v30"
|
||
|
||
@staticmethod
|
||
def generate_image(prompt, save_img_path):
|
||
"""生成图片"""
|
||
# 创建请求体
|
||
req = {
|
||
"req_key": JmTxt2Img.req_key,
|
||
"prompt": prompt
|
||
}
|
||
|
||
response_body = JmCommon.do_request("POST", {}, json.dumps(req).encode('utf-8'), JmTxt2Img.action)
|
||
jo = json.loads(response_body)
|
||
|
||
# 检查响应状态码
|
||
code = jo.get("code")
|
||
if not JmErrorCode.is_success(code):
|
||
error_msg = JmErrorCode.get_message_by_code(code)
|
||
print(f"生成图片失败: 错误码={code}, 错误信息={error_msg}")
|
||
raise Exception(f"生成图片失败: {error_msg}")
|
||
|
||
# 获取图片Base64数据
|
||
img_base64 = jo.get("data", {}).get("binary_data_base64", [""])[0]
|
||
|
||
# 确保目录存在
|
||
os.makedirs(os.path.dirname(save_img_path), exist_ok=True)
|
||
|
||
# 对 Base64 字符串进行解码并保存为文件
|
||
try:
|
||
# 注意：有些Base64字符串可能有前缀，需要去除
|
||
if img_base64.startswith('data:image'):
|
||
img_base64 = img_base64.split(',')[1]
|
||
bytes_data = base64.b64decode(img_base64)
|
||
with open(save_img_path, 'wb') as f:
|
||
f.write(bytes_data)
|
||
print(f"文件保存成功！文件位置: {save_img_path}")
|
||
except Exception as e:
|
||
print(f"保存图片失败: {str(e)}")
|
||
raise Exception(f"保存图片失败: {str(e)}")
|
||
|
||
@staticmethod
|
||
# 新增一个方法用于查询任务状态
|
||
@staticmethod
|
||
def query_task_status(task_id):
|
||
"""查询任务状态"""
|
||
req = {"task_id": task_id}
|
||
response_body = JmCommon.do_request("POST", {}, json.dumps(req).encode('utf-8'), "CVProcessStatus")
|
||
jo = json.loads(response_body)
|
||
|
||
code = jo.get("code")
|
||
if not JmErrorCode.is_success(code):
|
||
error_msg = JmErrorCode.get_message_by_code(code)
|
||
print(f"查询任务状态失败: 错误码={code}, 错误信息={error_msg}")
|
||
raise Exception(f"查询任务状态失败: {error_msg}")
|
||
|
||
return jo.get("data", {})
|
||
|
||
# 修改generate_image方法，使其支持异步模式
|
||
@staticmethod
|
||
def generate_image_async(prompt):
|
||
"""异步生成图片，返回任务ID"""
|
||
req = {"req_key": JmTxt2Img.req_key, "prompt": prompt}
|
||
response_body = JmCommon.do_request("POST", {}, json.dumps(req).encode('utf-8'), JmTxt2Img.action)
|
||
jo = json.loads(response_body)
|
||
|
||
code = jo.get("code")
|
||
if not JmErrorCode.is_success(code):
|
||
error_msg = JmErrorCode.get_message_by_code(code)
|
||
print(f"提交生成图片任务失败: 错误码={code}, 错误信息={error_msg}")
|
||
raise Exception(f"提交生成图片任务失败: {error_msg}")
|
||
|
||
return jo.get("data", {}).get("task_id")
|
||
|
||
# 修改main方法，实现异步轮询
|
||
@staticmethod
|
||
def main():
|
||
prompt = "鸿门宴,室内场景图..." # 详细提示词略
|
||
save_image_path = os.path.abspath("../../../../../Text2Img.jpg")
|
||
print(f"保存图片路径：{save_image_path}")
|
||
|
||
# 1. 提交生成任务
|
||
task_id = JmTxt2Img.generate_image_async(prompt)
|
||
print(f"图片生成任务已提交，任务ID: {task_id}")
|
||
|
||
# 2. 轮询任务状态
|
||
max_polling_times = 60 # 最多轮询60次
|
||
polling_interval = 5000 # 每5秒查询一次
|
||
polling_count = 0
|
||
|
||
while polling_count < max_polling_times:
|
||
try:
|
||
polling_count += 1
|
||
print(f"第{polling_count}次查询任务状态...")
|
||
task_status = JmTxt2Img.query_task_status(task_id)
|
||
|
||
# 检查任务是否完成
|
||
if task_status.get("status") == "completed":
|
||
print("图片生成完成！")
|
||
# 获取并保存图片
|
||
img_base64 = task_status.get("binary_data_base64", [""])[0]
|
||
os.makedirs(os.path.dirname(save_image_path), exist_ok=True)
|
||
|
||
if img_base64.startswith('data:image'):
|
||
img_base64 = img_base64.split(',')[1]
|
||
bytes_data = base64.b64decode(img_base64)
|
||
with open(save_image_path, 'wb') as f:
|
||
f.write(bytes_data)
|
||
print(f"文件保存成功！文件位置: {save_image_path}")
|
||
return
|
||
elif task_status.get("status") == "failed":
|
||
print(f"图片生成失败: {task_status.get('error_msg', '未知错误')}")
|
||
return
|
||
else:
|
||
print(f"任务进行中，当前状态: {task_status.get('status')}")
|
||
time.sleep(polling_interval / 1000)
|
||
except Exception as e:
|
||
print(f"查询任务状态异常: {str(e)}")
|
||
time.sleep(polling_interval / 1000)
|
||
|
||
print(f"已达到最大轮询次数: {max_polling_times}")
|
||
|
||
@staticmethod
|
||
def main():
|
||
"""主函数，用于测试"""
|
||
# 示例提示词
|
||
prompt = "鸿门宴,室内场景图。非真人，是一幅画。远景镜头，让人一看就知道是在讲鸿门宴，"
|
||
prompt += "1、要亮！"
|
||
prompt += "2、不要有人出现"
|
||
prompt += "3、有大厅、墙面，桌子,不要有台阶，否则项庄没机会刺杀刘邦"
|
||
prompt += "4、西周时期特点"
|
||
prompt += "5、宏大，雄伟,颜色鲜明,不要灰色调的，不要有雾霾之类的"
|
||
prompt += "6、超高清画质"
|
||
|
||
# 保存图片路径 - 使用绝对路径避免相对路径问题
|
||
save_image_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "..", "..", "Text2Img.jpg"))
|
||
print(f"保存图片路径：{save_image_path}")
|
||
|
||
# 添加重试逻辑，处理API并发限制错误
|
||
retry_count = 0
|
||
max_retries = 5 # 减少重试次数以便更快排查问题
|
||
retry_interval = 5000 # 重试间隔（毫秒）
|
||
|
||
while True:
|
||
try:
|
||
JmTxt2Img.generate_image(prompt, save_image_path)
|
||
# 成功生成图片，跳出循环
|
||
break
|
||
except Exception as e:
|
||
print(f"生成图片异常: {str(e)}")
|
||
retry_count += 1
|
||
if retry_count < max_retries:
|
||
print(f"等待{retry_interval}毫秒后重试...")
|
||
time.sleep(retry_interval / 1000)
|
||
else:
|
||
print(f"已达到最大重试次数: {max_retries}")
|
||
raise e # 达到最大重试次数，抛出异常
|
||
|
||
|
||
if __name__ == "__main__":
|
||
JmTxt2Img.main() |