'commit'

dsLightRag/JiMeng/JmTxt2Img.py

@@ -0,0 +1,90 @@
+import json
+import os
+import time
+import base64
+
+from JiMeng.Kit.JmCommon import JmCommon
+from JiMeng.Kit.JmErrorCode import JmErrorCode
+
+
+class JmTxt2Img:
+    action = "CVProcess"
+    req_key = "jimeng_high_aes_general_v21_L"
+
+    @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
+    def main():
+        """主函数，用于测试"""
+        # 示例提示词
+        prompt = "鸿门宴,室内场景图。非真人，是一幅画。远景镜头，让人一看就知道是在讲鸿门宴，"
+        prompt += "1、要亮！"
+        prompt += "2、不要有人出现"
+        prompt += "3、有大厅、墙面，桌子,不要有台阶，否则项庄没机会刺杀刘邦"
+        prompt += "4、西周时期特点"
+        prompt += "5、宏大，雄伟,颜色鲜明,不要灰色调的，不要有雾霾之类的"
+        prompt += "6、超高清画质"
+
+        # 保存图片路径
+        save_image_path = os.path.join(JmCommon.base_path, "Text2Img.jpg")
+        print(f"保存图片路径：{save_image_path}")
+
+        # 添加重试逻辑，处理API并发限制错误
+        retry_count = 0
+        max_retries = 1000  # 最大重试次数
+        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()