'commit'

dsLightRag/JiMeng/Kit/JmTxt2ImgUtil.py

@@ -0,0 +1,88 @@
+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}")