111 lines
4.5 KiB
Python
111 lines
4.5 KiB
Python
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 |