diff --git a/.gitignore b/.gitignore index 3c460d50..5035b170 100644 --- a/.gitignore +++ b/.gitignore @@ -13,4 +13,5 @@ /dsLightRag/.idea /dsLightRag/Topic/ /dsLightRag/Dsideal/YunXiao/Down/ +/dsLightRag/.venv/ diff --git a/dsLightRag/JiMeng/Async_Txt2Img.py b/dsLightRag/JiMeng/Async_Txt2Img.py new file mode 100644 index 00000000..f9a903b4 --- /dev/null +++ b/dsLightRag/JiMeng/Async_Txt2Img.py @@ -0,0 +1,14 @@ +# pip install --user volcengine +from JiMeng.JmTxt2Img import JmTxt2Img +from JiMeng.Kit.JmCommon import JmCommon + +if __name__ == '__main__': + visual_service = VisualService() + + visual_service.set_ak(JmCommon.ak) + visual_service.set_sk(JmCommon.sk) + + prompt = "鸿门宴,室内场景图..." + req = {"req_key": JmTxt2Img.req_key, "prompt": prompt} + resp = visual_service.cv_submit_task(req) + print(resp) \ No newline at end of file diff --git a/dsLightRag/JiMeng/JmTxt2Img.py b/dsLightRag/JiMeng/JmTxt2Img.py index 61a6a8ce..b34fda3f 100644 --- a/dsLightRag/JiMeng/JmTxt2Img.py +++ b/dsLightRag/JiMeng/JmTxt2Img.py @@ -52,6 +52,87 @@ class JmTxt2Img: print(f"保存图片失败: {str(e)}") raise Exception(f"保存图片失败: {str(e)}") + @staticmethod + # 新增一个方法用于查询任务状态 + @staticmethod + def query_task_status(task_id): + """查询任务状态""" + req = {"task_id": task_id} + response_body = JmCommon.do_request("POST", {}, json.dumps(req).encode('utf-8'), "CVProcessStatus") + 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}") + + return jo.get("data", {}) + + # 修改generate_image方法,使其支持异步模式 + @staticmethod + def generate_image_async(prompt): + """异步生成图片,返回任务ID""" + 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}") + + return jo.get("data", {}).get("task_id") + + # 修改main方法,实现异步轮询 + @staticmethod + def main(): + prompt = "鸿门宴,室内场景图..." # 详细提示词略 + save_image_path = os.path.abspath("../../../../../Text2Img.jpg") + print(f"保存图片路径:{save_image_path}") + + # 1. 提交生成任务 + task_id = JmTxt2Img.generate_image_async(prompt) + print(f"图片生成任务已提交,任务ID: {task_id}") + + # 2. 轮询任务状态 + max_polling_times = 60 # 最多轮询60次 + polling_interval = 5000 # 每5秒查询一次 + polling_count = 0 + + while polling_count < max_polling_times: + try: + polling_count += 1 + print(f"第{polling_count}次查询任务状态...") + task_status = JmTxt2Img.query_task_status(task_id) + + # 检查任务是否完成 + if task_status.get("status") == "completed": + print("图片生成完成!") + # 获取并保存图片 + img_base64 = task_status.get("binary_data_base64", [""])[0] + os.makedirs(os.path.dirname(save_image_path), exist_ok=True) + + if img_base64.startswith('data:image'): + img_base64 = img_base64.split(',')[1] + bytes_data = base64.b64decode(img_base64) + with open(save_image_path, 'wb') as f: + f.write(bytes_data) + print(f"文件保存成功!文件位置: {save_image_path}") + return + elif task_status.get("status") == "failed": + print(f"图片生成失败: {task_status.get('error_msg', '未知错误')}") + return + else: + print(f"任务进行中,当前状态: {task_status.get('status')}") + time.sleep(polling_interval / 1000) + except Exception as e: + print(f"查询任务状态异常: {str(e)}") + time.sleep(polling_interval / 1000) + + print(f"已达到最大轮询次数: {max_polling_times}") + @staticmethod def main(): """主函数,用于测试""" diff --git a/dsLightRag/JiMeng/__pycache__/JmTxt2Img.cpython-310.pyc b/dsLightRag/JiMeng/__pycache__/JmTxt2Img.cpython-310.pyc new file mode 100644 index 00000000..c6b613e4 Binary files /dev/null and b/dsLightRag/JiMeng/__pycache__/JmTxt2Img.cpython-310.pyc differ