Files
dsProject/dsLightRag/JiMeng/Kit/JmTxt2ImgUtil.py
2025-08-21 09:18:06 +08:00

89 lines
3.2 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import base64
import json
import logging
import uuid
from Config.Config import OBS_SERVER, OBS_PREFIX, OBS_BUCKET
from JiMeng.Kit.JmCommon import JmCommon
from JiMeng.Kit.JmErrorCode import JmErrorCode
from Util.ObsUtil import ObsUploader
# 设置日志
logger = logging.getLogger(__name__)
logger.level = logging.INFO
# 添加控制台处理器
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)
# 设置日志格式
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
console_handler.setFormatter(formatter)
# 将处理器添加到logger
logger.addHandler(console_handler)
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):
"""生成图片"""
# 创建请求体
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]
try:
# 注意有些Base64字符串可能有前缀需要去除
if img_base64.startswith('data:image'):
img_base64 = img_base64.split(',')[1]
bytes_data = base64.b64decode(img_base64)
# 创建上传器实例
uploader = ObsUploader()
# 生成一个随机的图片文件名
jpg_file_name = f"{str(uuid.uuid4())}.jpg"
# 上传参数
object_key = OBS_PREFIX + "/JiMeng/" + jpg_file_name
# 执行上传
success, result = uploader.upload_base64_image(object_key, bytes_data)
# 处理结果
if success:
logger.info(f'{jpg_file_name}上传成功!')
# 获取上传文件的 URL
file_url = f"https://{OBS_BUCKET}.{OBS_SERVER}/{object_key}"
return file_url
except Exception as e:
print(f"保存图片失败: {str(e)}")
raise Exception(f"保存图片失败: {str(e)}")
return None
if __name__ == "__main__":
prompt = "雨天街头,一位头戴黄色安全帽、身穿荧光黄安全背心与棕色外套、搭配蓝色牛仔裤和棕色工装靴的中国工人,背着一个中国小女孩。工人脖子上侧挂一个小熊玩偶包,拖住小女孩臀部的手同时挂着一袋小熊饼干。小女孩双手搂着工人的脖子。小女孩穿红色雨衣,戴红色雨衣帽。两个人露出微笑。背景是居民楼和树木,地面潮湿,氛围温馨。"
# 保存图片路径
url = JmTxt2Img.generate_image(prompt)
logger.info(f"图片URL: {url}")