'commit'

dsLightRag/Liblib/Kit/PuLIDGenerator.py

@@ -0,0 +1,111 @@
+import json
+
+from Config.Config import OBS_BUCKET, OBS_SERVER
+from LibLibGenerator import LibLibGenerator
+
+class PuLIDGenerator:
+    def __init__(self):
+        self.liblib_util = LibLibGenerator()
+        # 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