dsProject/dsLightRag/Liblib/PuLIDGenerator.py

import json

from Config.Config import OBS_BUCKET, OBS_SERVER
from LibLibUtil import LibLibUtil

class PuLIDGenerator:
    def __init__(self):
        self.liblib_util = LibLibUtil()
        # PuLID人像换脸参数示例
        self.default_params = {
            "templateUuid": "6f7c4652458d4802969f8d089cf5b91f",
            "steps": 20,
            "width": 768,
            "height": 1024,
            "imgCount": 1
        }

    def set_template_uuid(self, template_uuid):
        """设置模板UUID"""
        self.template_uuid = template_uuid
        return self

    def set_default_params(self, **kwargs):
        """设置默认生成参数"""
        self.default_params.update(kwargs)
        return self

    def build_control_net_params(self, reference_image_url, control_weight=1.0):
        """构建控制网络参数"""
        return [{
            "unitOrder": 0,
            "sourceImage": reference_image_url,
            "width": self.default_params["width"],
            "height": self.default_params["height"],
            "preprocessor": 0,
            "annotationParameters": {"none": {}},
            "model": "405836d1ae2646b4ba2716ed6bd5453a",
            "controlWeight": control_weight,
            "startingControlStep": 0,
            "endingControlStep": 1,
            "pixelPerfect": 1,
            "controlMode": 0,
            "resizeMode": 1
        }]

    def generate_image(self, prompt, reference_image_url, control_weight=1.0):
        """根据提示词和参考图片生成图像，并自动处理后续流程"""
        # 构建生成参数
        generate_params = {
            "templateUuid":self.default_params["templateUuid"],
            "generateParams": {
                "prompt": prompt,
                **self.default_params,
                "controlNet": self.build_control_net_params(reference_image_url, control_weight)
            }
        }

        # 调用生成接口
        try:
            response = self.liblib_util.post_request(
                "/api/generate/webui/text2img",
                generate_params
            )

            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: {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}")
                            file_url = f"https://{OBS_BUCKET}.{OBS_SERVER}/{obs_url}"
                            return file_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}")
                return None
        except Exception as e:
            print(f"请求发生异常: {str(e)}")
            return None