dsProject/dsLightRag/JiMeng/T1_JmTxt2Img.py

import base64
import json
import os

from Config.Config import OBS_SERVER, OBS_PREFIX
from JiMeng.Kit.JmCommon import JmCommon
from JiMeng.Kit.JmErrorCode import JmErrorCode
from Util.ObsUtil import ObsUploader


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}")

            # 创建上传器实例
            uploader = ObsUploader()
            # 上传参数
            object_key = OBS_PREFIX + "/" + file_path
            # 执行上传
            success, result = uploader.upload_file(object_key, file_path)
            # 处理结果
            if success:
                logger.info(f'{file_path}上传成功!')
                # 获取上传文件的 URL
                file_url = f"https://{OBS_BUCKET}.{OBS_SERVER}/{object_key}"
                img_url_list.append(file_url)

        except Exception as e:
            print(f"保存图片失败: {str(e)}")
            raise Exception(f"保存图片失败: {str(e)}")



    @staticmethod
    def main():
        # 示例提示词
        prompt = "雨天街头，一位头戴黄色安全帽、身穿荧光黄安全背心与棕色外套、搭配蓝色牛仔裤和棕色工装靴的中国工人，背着一个中国小女孩。工人脖子上侧挂一个小熊玩偶包，拖住小女孩臀部的手同时挂着一袋小熊饼干。小女孩双手搂着工人的脖子。小女孩穿红色雨衣，戴红色雨衣帽。两个人露出微笑。背景是居民楼和树木，地面潮湿，氛围温馨。"

        # 保存图片路径
        save_image_path =  r"D:\Text2Img.jpg"
        JmTxt2Img.generate_image(prompt, save_image_path)
        print(f"保存图片路径：{save_image_path}")


if __name__ == "__main__":
    JmTxt2Img.main()