'commit'
This commit is contained in:
304
dsLightRag/Liblib/Kit/LibLibGenerator.py
Normal file
304
dsLightRag/Liblib/Kit/LibLibGenerator.py
Normal file
@@ -0,0 +1,304 @@
|
||||
import hmac
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from hashlib import sha1
|
||||
import base64
|
||||
import time
|
||||
import uuid
|
||||
import requests
|
||||
import Config.Config
|
||||
from Util.ObsUtil import ObsUploader
|
||||
|
||||
|
||||
class LibLibGenerator:
|
||||
def __init__(self):
|
||||
self.base_url = Config.Config.LIBLIB_URL
|
||||
self.access_key = Config.Config.LIBLIB_ACCESSKEY
|
||||
self.secret_key = Config.Config.LIBLIB_SECRETKEY
|
||||
self.timeout = 30
|
||||
|
||||
def make_sign(self, uri):
|
||||
"""生成API请求签名"""
|
||||
timestamp = str(int(time.time() * 1000))
|
||||
signature_nonce = str(uuid.uuid4())
|
||||
content = '&'.join((uri, timestamp, signature_nonce))
|
||||
|
||||
digest = hmac.new(self.secret_key.encode(), content.encode(), sha1).digest()
|
||||
sign = base64.urlsafe_b64encode(digest).rstrip(b'=').decode()
|
||||
return sign, timestamp, signature_nonce
|
||||
|
||||
def post_request(self, uri, payload):
|
||||
"""发送POST请求到Liblib API"""
|
||||
try:
|
||||
sign, timestamp, signature_nonce = self.make_sign(uri)
|
||||
url = f'{self.base_url}{uri}?AccessKey={self.access_key}&Signature={sign}&Timestamp={timestamp}&SignatureNonce={signature_nonce}'
|
||||
|
||||
headers = {'Content-Type': 'application/json'}
|
||||
response = requests.post(url, json=payload, headers=headers, timeout=self.timeout)
|
||||
response.raise_for_status()
|
||||
|
||||
response_data = response.json()
|
||||
if response_data.get('code') == 0:
|
||||
return response_data.get('data')
|
||||
else:
|
||||
print(f"url={url}")
|
||||
print(f"response_data={response_data}")
|
||||
print(f"API错误: {response_data.get('msg')}")
|
||||
# print(response_data)
|
||||
return None
|
||||
except requests.exceptions.RequestException as e:
|
||||
print(f"请求异常: {str(e)}")
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f"处理异常: {str(e)}")
|
||||
return None
|
||||
|
||||
def get_model_version_info(self, version_uuid):
|
||||
"""获取模型版本信息"""
|
||||
uri = "/api/model/version/get"
|
||||
payload = {"versionUuid": version_uuid}
|
||||
|
||||
model_info = self.post_request(uri, payload)
|
||||
if model_info:
|
||||
return {
|
||||
'modelName': model_info.get('modelName'),
|
||||
'versionName': model_info.get('versionName'),
|
||||
'commercialUse': model_info.get('commercialUse'),
|
||||
'modelUrl': model_info.get('modelUrl')
|
||||
}
|
||||
else:
|
||||
print("获取模型版本信息失败")
|
||||
print(model_info)
|
||||
return None
|
||||
|
||||
|
||||
def get_generation_status(self, generate_uuid):
|
||||
"""查询生图任务状态和结果"""
|
||||
uri = "/api/generate/webui/status"
|
||||
payload = {"generateUuid": generate_uuid}
|
||||
return self.post_request(uri, payload)
|
||||
|
||||
def wait_for_generation_completion(self, generate_uuid, interval=5, max_wait_time=100):
|
||||
"""等待生成完成并返回最终状态"""
|
||||
while True:
|
||||
status_data = self.get_generation_status(generate_uuid)
|
||||
if not status_data:
|
||||
print("获取状态失败，重试中...")
|
||||
time.sleep(interval)
|
||||
continue
|
||||
|
||||
# 获取当前状态和进度
|
||||
generate_status = status_data.get('generateStatus')
|
||||
percent = status_data.get('percentCompleted', 0) * 100
|
||||
status_text = {
|
||||
0: "初始化",
|
||||
1: "处理中",
|
||||
2: "排队中",
|
||||
3: "已取消",
|
||||
4: "失败",
|
||||
5: "完成"
|
||||
}.get(generate_status, f"未知状态({generate_status})")
|
||||
|
||||
# 显示进度和状态
|
||||
print(f"生成进度：{percent:.2f}% | 状态：{status_text}")
|
||||
|
||||
# 检查是否完成
|
||||
if generate_status in [3, 4, 5]: # 取消、失败或完成
|
||||
return status_data
|
||||
|
||||
time.sleep(interval)
|
||||
|
||||
print(f"生图任务超时（{max_wait_time}秒）")
|
||||
return None
|
||||
|
||||
def download_and_upload_to_obs(slef, image_url, generate_uuid):
|
||||
"""下载图片并上传到OBS"""
|
||||
try:
|
||||
# 1. 清理URL（去除可能的引号和空格）
|
||||
clean_url = image_url.strip('` ')
|
||||
print(f"清理后的图片URL: {clean_url}")
|
||||
|
||||
# 2. 下载图片到临时文件
|
||||
response = requests.get(clean_url, stream=True)
|
||||
response.raise_for_status() # 检查HTTP错误
|
||||
|
||||
# 创建临时文件
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_file:
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
temp_file.write(chunk)
|
||||
temp_file_path = temp_file.name
|
||||
print(f"图片已下载至临时文件: {temp_file_path}")
|
||||
|
||||
# 3. 上传到OBS
|
||||
obs_uploader = ObsUploader()
|
||||
obs_object_key = f"HuangHai/LibLib/{generate_uuid}.jpg"
|
||||
success, result = obs_uploader.upload_file(
|
||||
object_key=obs_object_key,
|
||||
file_path=temp_file_path
|
||||
)
|
||||
|
||||
# 4. 清理临时文件
|
||||
os.unlink(temp_file_path)
|
||||
|
||||
if success:
|
||||
print(f"✅ 图片成功上传至OBS")
|
||||
return obs_object_key
|
||||
else:
|
||||
print(f"❌ OBS上传失败: {result}")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
print(f"处理图片时发生错误: {str(e)}")
|
||||
return None
|
||||
|
||||
def generate_default_text_to_image(self, prompt, steps=20, width=768, height=1024,
|
||||
img_count=1, seed=-1, restore_faces=0,
|
||||
additional_network=None):
|
||||
"""
|
||||
使用默认官方模型进行文生图
|
||||
- Checkpoint默认为官方模型
|
||||
- 可用模型范围：基础算法F.1
|
||||
- 支持additional network
|
||||
"""
|
||||
generate_params = {
|
||||
"templateUuid": "6f7c4652458d4802969f8d089cf5b91f", # 参数模板ID F.1文生图
|
||||
"generateParams": {
|
||||
"prompt": prompt,
|
||||
"steps": steps,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"imgCount": img_count,
|
||||
"seed": seed,
|
||||
"restoreFaces": restore_faces
|
||||
}
|
||||
}
|
||||
|
||||
# 添加additional network参数（如果有）
|
||||
if additional_network:
|
||||
generate_params["generateParams"]["additionalNetwork"] = additional_network
|
||||
|
||||
# 调用生成接口
|
||||
response = self.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.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.download_and_upload_to_obs(image_url, generate_uuid)
|
||||
|
||||
if obs_url:
|
||||
print(f"✅ 文件处理完成，OBS地址: {obs_url}")
|
||||
file_url = f"https://{Config.Config.OBS_BUCKET}.{Config.Config.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
|
||||
|
||||
|
||||
def generate_custom_checkpoint_text_to_image(self, template_uuid, checkpoint_id,
|
||||
prompt="", negative_prompt="", steps=20,
|
||||
sampler=15, cfg_scale=7, width=768,
|
||||
height=1024, img_count=1, randn_source=0,
|
||||
seed=-1, restore_faces=0, hi_res_fix_info=None):
|
||||
"""
|
||||
使用自定义Checkpoint模型进行文生图
|
||||
"""
|
||||
generate_params = {
|
||||
"templateUuid": template_uuid,
|
||||
"generateParams": {
|
||||
"checkPointId": checkpoint_id,
|
||||
"prompt": prompt,
|
||||
"negativePrompt": negative_prompt,
|
||||
"sampler": sampler,
|
||||
"steps": steps,
|
||||
"cfgScale": cfg_scale,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"imgCount": img_count,
|
||||
"randnSource": randn_source,
|
||||
"seed": seed,
|
||||
"restoreFaces": restore_faces
|
||||
}
|
||||
}
|
||||
|
||||
# 添加高分辨率修复参数（如果有）
|
||||
if hi_res_fix_info:
|
||||
generate_params["generateParams"]["hiResFixInfo"] = hi_res_fix_info
|
||||
|
||||
# 调用生成接口
|
||||
response = self.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.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.download_and_upload_to_obs(image_url, generate_uuid)
|
||||
|
||||
if obs_url:
|
||||
print(f"✅ 文件处理完成，OBS地址: {obs_url}")
|
||||
file_url = f"https://{Config.Config.OBS_BUCKET}.{Config.Config.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
|
||||
32
dsLightRag/Liblib/Kit/LibLibSamplingEnum.py
Normal file
32
dsLightRag/Liblib/Kit/LibLibSamplingEnum.py
Normal file
@@ -0,0 +1,32 @@
|
||||
|
||||
# 采样方法字典 - 包含28种采样方法及其枚举值和推荐度
|
||||
SAMPLING_METHODS = {
|
||||
"Euler a": 0,
|
||||
"Euler": 1,
|
||||
"LMS": 2,
|
||||
"Heun": 3,
|
||||
"DPM2": 4,
|
||||
"DPM2 a": 5,
|
||||
"DPM++ 2S a": 6,
|
||||
"DPM++ 2M": 7,
|
||||
"DPM++ SDE": 8,
|
||||
"DPM fast": 9,
|
||||
"DPM adaptive": 10,
|
||||
"LMS Karras": 11,
|
||||
"DPM2 Karras": 12,
|
||||
"DPM2 a Karras": 13,
|
||||
"DPM++ 2S a Karras": 14,
|
||||
"DPM++ 2M Karras": 15,
|
||||
"DPM++ SDE Karras": 16,
|
||||
"DDIM": 17,
|
||||
"PLMS": 18,
|
||||
"UNIPC": 19,
|
||||
"DPM++ 2M SDE Karras": 20,
|
||||
"DPM++ 2M SDE EXPONENTIAL": 21,
|
||||
"DPM++ 2M SDE Heun Karras": 24,
|
||||
"DPM++ 2M SDE Heun Exponential": 25,
|
||||
"DPM++ 3M SDE Karras": 27,
|
||||
"DPM++ 3M SDE Exponential": 28,
|
||||
"Restart": 29,
|
||||
"LCM": 30
|
||||
}
|
||||
111
dsLightRag/Liblib/Kit/PuLIDGenerator.py
Normal file
111
dsLightRag/Liblib/Kit/PuLIDGenerator.py
Normal file
@@ -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
|
||||
0
dsLightRag/Liblib/Kit/__init__.py
Normal file
0
dsLightRag/Liblib/Kit/__init__.py
Normal file
Reference in New Issue
Block a user