diff --git a/dsLightRag/Liblib/LiblibUtil.py b/dsLightRag/Liblib/LibLibUtil.py similarity index 76% rename from dsLightRag/Liblib/LiblibUtil.py rename to dsLightRag/Liblib/LibLibUtil.py index 34c5aea1..9f8fe6b8 100644 --- a/dsLightRag/Liblib/LiblibUtil.py +++ b/dsLightRag/Liblib/LibLibUtil.py @@ -1,4 +1,5 @@ import hmac +import json import os import tempfile from hashlib import sha1 @@ -10,7 +11,7 @@ import Config.Config from Util.ObsUtil import ObsUploader -class LiblibUtil: +class LibLibUtil: def __init__(self): self.base_url = Config.Config.LIBLIB_URL self.access_key = Config.Config.LIBLIB_ACCESSKEY @@ -50,7 +51,6 @@ class LiblibUtil: print(f"处理异常: {str(e)}") return None - # ===== 从T1_VersionGet.py迁移的方法 ===== def get_model_version_info(self, version_uuid): """获取模型版本信息""" uri = "/api/model/version/get" @@ -75,7 +75,6 @@ class LiblibUtil: } return self.post_request(uri, payload) - # ===== 从T2.py迁移的方法 ===== def get_generation_status(self, generate_uuid): """查询生图任务状态和结果""" uri = "/api/generate/webui/status" @@ -115,7 +114,6 @@ class LiblibUtil: print(f"生图任务超时（{max_wait_time}秒）") return None - @staticmethod def download_and_upload_to_obs(image_url, generate_uuid): """下载图片并上传到OBS""" try: @@ -155,3 +153,35 @@ class LiblibUtil: except Exception as e: print(f"处理图片时发生错误: {str(e)}") return None + + def process_generation_task(self, generate_uuid, interval=2): + """ + 完整处理生成任务: 轮询状态 → 下载图片 → 上传OBS → 清理临时文件 + :param generate_uuid: 生成任务UUID + :param interval: 轮询间隔(秒) + :return: OBS路径或None + """ + try: + print(f"开始监控生成任务 {generate_uuid}，每{interval}秒检查一次...") + + # 轮询等待生成完成 + status_data = self.wait_for_generation_completion(generate_uuid, interval) + print(f"生图状态: {json.dumps(status_data, ensure_ascii=False, indent=2)}") + + # 检查生成状态是否为5（完成） + if status_data and status_data.get('generateStatus') == 5: + # 提取图片URL + if status_data.get('images') and len(status_data['images']) > 0: + image_url = status_data['images'][0]['imageUrl'] + # 下载并上传到OBS + obs_path = self.download_and_upload_to_obs(image_url, generate_uuid) + return obs_path + else: + print("❌ 未找到图片数据") + return None + else: + print(f"❌ 生图未完成或失败，状态码: {status_data.get('generateStatus')}") + return None + except Exception as e: + print(f"处理生成任务时发生异常: {str(e)}") + return None diff --git a/dsLightRag/Liblib/ControlNetGenerator.py b/dsLightRag/Liblib/PuLIDGenerator.py similarity index 54% rename from dsLightRag/Liblib/ControlNetGenerator.py rename to dsLightRag/Liblib/PuLIDGenerator.py index fa4d9315..a49c6aa3 100644 --- a/dsLightRag/Liblib/ControlNetGenerator.py +++ b/dsLightRag/Liblib/PuLIDGenerator.py @@ -1,9 +1,9 @@ import json -from LiblibUtil import LiblibUtil +from LibLibUtil import LibLibUtil -class ControlNetImageGenerator: +class PuLIDGenerator: def __init__(self): - self.liblib_util = LiblibUtil() + self.liblib_util = LibLibUtil() # PuLID人像换脸参数示例 self.template_uuid = "6f7c4652458d4802969f8d089cf5b91f" self.default_params = { @@ -42,7 +42,7 @@ class ControlNetImageGenerator: }] def generate_image(self, prompt, reference_image_url, control_weight=1.0): - """根据提示词和参考图片生成图像""" + """根据提示词和参考图片生成图像，并自动处理后续流程""" # 构建生成参数 generate_params = { "templateUuid": self.template_uuid, @@ -63,10 +63,42 @@ class ControlNetImageGenerator: print(f"API响应: {json.dumps(response, ensure_ascii=False, indent=2)}") if response and "generateUuid" in response: + generate_uuid = response["generateUuid"] print("✅ 图像生成任务已成功提交！") - print(f"生成UUID: {response['generateUuid']}") - print("提示: 可使用T2.py中的get_generation_status方法查询生成进度") - return response + print(f"生成UUID: {generate_uuid}") + print("开始轮询生成状态...") + + # 每2秒探测一次生成状态，直到完成 + status_data = self.liblib_util.wait_for_generation_completion( + generate_uuid, interval=2 + ) + + # 检查生成状态 + if status_data and status_data.get("generateStatus") == 5: + print("🎉 图像生成完成！开始处理文件...") + + # 提取图片URL + if status_data.get("images") and len(status_data["images"]) > 0: + image_url = status_data["images"][0]["imageUrl"] + + # 下载并上传到OBS + obs_url = self.liblib_util.download_and_upload_to_obs( + image_url, generate_uuid + ) + + if obs_url: + print(f"✅ 文件处理完成，OBS地址: {obs_url}") + return obs_url + else: + print("❌ 文件上传OBS失败") + return None + else: + print("❌ 未找到生成的图片数据") + return None + else: + error_msg = status_data.get('message', '未知错误') if status_data else '生成状态查询失败' + print(f"❌ 图像生成失败: {error_msg}") + return None else: error_msg = response.get('message', '未知错误') if response else 'API无响应' print(f"❌ 图像生成失败: {error_msg}") diff --git a/dsLightRag/Liblib/T1.py b/dsLightRag/Liblib/T1.py index 8b34d5b8..f8783fdd 100644 --- a/dsLightRag/Liblib/T1.py +++ b/dsLightRag/Liblib/T1.py @@ -1,8 +1,8 @@ -from .LiblibUtil import LiblibUtil +from .LibLibUtil import LibLibUtil if __name__ == '__main__': - liblib = LiblibUtil() + liblib = LibLibUtil() # 测试获取模型版本信息 print("===== 测试获取模型版本信息 =====") diff --git a/dsLightRag/Liblib/T2.py b/dsLightRag/Liblib/T2.py index 9803ed41..c956cc94 100644 --- a/dsLightRag/Liblib/T2.py +++ b/dsLightRag/Liblib/T2.py @@ -1,30 +1,16 @@ -import json - -from LiblibUtil import LiblibUtil +from LibLibUtil import LibLibUtil if __name__ == '__main__': - liblib = LiblibUtil() - print("===== 测试查询生图结果并上传OBS =====") + liblib = LibLibUtil() + print("===== 测试生成任务处理流程 =====") test_generate_uuid = "df9bd6e03a204542b1daaba2f17c42e6" if test_generate_uuid: - # 等待任务完成 - status_data = liblib.wait_for_generation_completion(test_generate_uuid) - print(f"生图状态: {json.dumps(status_data, ensure_ascii=False, indent=2)}") - - # 检查生成状态是否为5（完成） - if status_data and status_data.get('generateStatus') == 5: - # 提取图片URL - if status_data.get('images') and len(status_data['images']) > 0: - image_url = status_data['images'][0]['imageUrl'] - # 下载并上传到OBS - # 使用实例调用而非类调用 - obs_path = liblib.download_and_upload_to_obs(image_url, test_generate_uuid) - if obs_path: - print(f"\n最终OBS路径: {obs_path}") - else: - print("❌ 未找到图片数据") + # 调用封装好的完整处理方法 + obs_url = liblib.process_generation_task(test_generate_uuid, interval=2) + if obs_url: + print(f"\n✅ 处理完成，OBS地址: {obs_url}") else: - print(f"❌ 生图未完成或失败，状态码: {status_data.get('generateStatus')}") + print("\n❌ 任务处理失败") else: print("请先设置有效的生图任务UUID") \ No newline at end of file diff --git a/dsLightRag/Liblib/T3.py b/dsLightRag/Liblib/T3.py index f7ba6bda..8fe1a9ce 100644 --- a/dsLightRag/Liblib/T3.py +++ b/dsLightRag/Liblib/T3.py @@ -1,17 +1,15 @@ -import json -from LiblibUtil import LiblibUtil -from ControlNetGenerator import ControlNetImageGenerator +from PuLIDGenerator import PuLIDGenerator if __name__ == "__main__": # 创建生成器实例 - generator = ControlNetImageGenerator() - + generator = PuLIDGenerator() # 可选：修改默认参数 generator.set_default_params(steps=25, width=800, height=1000) - # 生成图像 - generator.generate_image( + obs_url = generator.generate_image( prompt="filmfotos, Asian portrait,A young woman wearing a green baseball cap,covering one eye with her hand", reference_image_url="https://dsideal.obs.cn-north-1.myhuaweicloud.com/HuangHai/Backup/HuangWanQiao.jpg", control_weight=0.8 - ) \ No newline at end of file + ) + if obs_url: + print(f"最终OBS地址: {obs_url}") diff --git a/dsLightRag/Liblib/__pycache__/LiblibUtil.cpython-310.pyc b/dsLightRag/Liblib/__pycache__/LiblibUtil.cpython-310.pyc deleted file mode 100644 index 116420cf..00000000 Binary files a/dsLightRag/Liblib/__pycache__/LiblibUtil.cpython-310.pyc and /dev/null differ