'commit'

dsLightRag/JiMeng/T1_JmTxt2Img.py

@@ -1,22 +1,39 @@
 import base64
 import json
+import logging
 import os
+import uuid
 
-from Config.Config import OBS_SERVER, OBS_PREFIX
+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, save_img_path):
+    def generate_image(prompt):
         """生成图片"""
         # 创建请求体
         req = {
@@ -37,48 +54,40 @@ class JmTxt2Img:
         # 获取图片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()
+
+            # 生成一个随机的图片文件名
+            jpg_file_name = f"{str(uuid.uuid4())}.jpg"
             # 上传参数
-            object_key = OBS_PREFIX + "/" + file_path
+            object_key = OBS_PREFIX + "/JiMeng/" + jpg_file_name
             # 执行上传
-            success, result = uploader.upload_file(object_key, file_path)
+            success, result = uploader.upload_base64_image(object_key, bytes_data)
             # 处理结果
             if success:
-                logger.info(f'{file_path}上传成功!')
+                logger.info(f'{jpg_file_name}上传成功!')
                 # 获取上传文件的 URL
                 file_url = f"https://{OBS_BUCKET}.{OBS_SERVER}/{object_key}"
-                img_url_list.append(file_url)
+                return file_url
 
         except Exception as e:
             print(f"保存图片失败: {str(e)}")
             raise Exception(f"保存图片失败: {str(e)}")
-
-
+        return None
 
     @staticmethod
     def main():
-        # 示例提示词
         prompt = "雨天街头，一位头戴黄色安全帽、身穿荧光黄安全背心与棕色外套、搭配蓝色牛仔裤和棕色工装靴的中国工人，背着一个中国小女孩。工人脖子上侧挂一个小熊玩偶包，拖住小女孩臀部的手同时挂着一袋小熊饼干。小女孩双手搂着工人的脖子。小女孩穿红色雨衣，戴红色雨衣帽。两个人露出微笑。背景是居民楼和树木，地面潮湿，氛围温馨。"
-
         # 保存图片路径
-        save_image_path =  r"D:\Text2Img.jpg"
-        JmTxt2Img.generate_image(prompt, save_image_path)
-        print(f"保存图片路径：{save_image_path}")
+        url = JmTxt2Img.generate_image(prompt)
+        logger.info(f"图片URL: {url}")
 
 
 if __name__ == "__main__":
-    JmTxt2Img.main()
+    JmTxt2Img.main()