140 lines
4.5 KiB
Python
140 lines
4.5 KiB
Python
import requests
|
||
import uuid
|
||
import json
|
||
import os
|
||
import time
|
||
from PIL import Image
|
||
from io import BytesIO
|
||
|
||
|
||
class ComfyUIClient:
|
||
def __init__(self, server_url="http://127.0.0.1:8188"):
|
||
"""初始化ComfyUI客户端"""
|
||
self.server_url = server_url.rstrip("/")
|
||
self.client_id = str(uuid.uuid4())
|
||
self.output_images = []
|
||
|
||
def upload_image(self, image_path):
|
||
"""上传图片到ComfyUI服务器"""
|
||
if not os.path.exists(image_path):
|
||
raise FileNotFoundError(f"图片文件不存在: {image_path}")
|
||
|
||
# 读取图片文件
|
||
with open(image_path, "rb") as file:
|
||
image_data = file.read()
|
||
|
||
filename = os.path.basename(image_path)
|
||
files = {
|
||
"image": (filename, image_data, "image/png")
|
||
}
|
||
data = {
|
||
"overwrite": "true"
|
||
}
|
||
|
||
url = f"{self.server_url}/upload/image"
|
||
response = requests.post(url, files=files, data=data)
|
||
|
||
if response.status_code != 200:
|
||
raise Exception(f"上传图片失败: {response.text}")
|
||
|
||
print(f"成功上传图片: {filename}")
|
||
return filename
|
||
|
||
def submit_prompt(self, workflow, input_image_path):
|
||
"""提交绘图任务"""
|
||
# 上传图片并修改工作流
|
||
image_filename = self.upload_image(input_image_path)
|
||
|
||
# 修改工作流中的图片输入
|
||
if "10" in workflow:
|
||
workflow["10"]["inputs"]["image"] = image_filename
|
||
|
||
# 构建请求数据
|
||
data = {
|
||
"client_id": self.client_id,
|
||
"prompt": workflow
|
||
}
|
||
|
||
# 提交任务
|
||
url = f"{self.server_url}/prompt"
|
||
response = requests.post(url, json=data)
|
||
|
||
if response.status_code != 200:
|
||
raise Exception(f"提交任务失败: {response.text}")
|
||
|
||
result = response.json()
|
||
prompt_id = result.get("prompt_id")
|
||
print(f"任务提交成功，任务ID: {prompt_id}")
|
||
return prompt_id
|
||
|
||
def wait_for_completion(self, prompt_id, timeout=300, interval=2):
|
||
"""通过 /history 轮询等待任务完成"""
|
||
history_url = f"{self.server_url}/history/{prompt_id}"
|
||
start_time = time.time()
|
||
|
||
while time.time() - start_time < timeout:
|
||
try:
|
||
response = requests.get(history_url)
|
||
if response.status_code == 200:
|
||
history = response.json().get(prompt_id)
|
||
if history and history.get("status", {}).get("completed", False):
|
||
outputs = history.get("outputs", {})
|
||
for node_id, output in outputs.items():
|
||
if "images" in output:
|
||
self.output_images = output["images"]
|
||
print(
|
||
f"任务完成，生成了 {len(self.output_images)} 张图片")
|
||
return self.output_images
|
||
except Exception as e:
|
||
print(f"轮询错误: {e}")
|
||
|
||
time.sleep(interval)
|
||
|
||
print("任务超时或未生成图片")
|
||
return None
|
||
|
||
def get_image(self, filename, save_path=None):
|
||
"""获取并可选保存生成的图片"""
|
||
print(f"获取图片名称: {filename}")
|
||
url = f"{self.server_url}/view"
|
||
params = {
|
||
"filename": filename,
|
||
"type": "temp" # 这里可以根据需要修改为 "output" 或 "temp"
|
||
}
|
||
response = requests.get(url, params=params)
|
||
|
||
if response.status_code != 200:
|
||
raise Exception(f"获取图片失败: {response.status_code}")
|
||
|
||
image = Image.open(BytesIO(response.content))
|
||
|
||
if save_path:
|
||
os.makedirs(os.path.dirname(
|
||
os.path.abspath(save_path)), exist_ok=True)
|
||
image.save(save_path)
|
||
print(f"图片已保存到: {save_path}")
|
||
|
||
return image
|
||
|
||
|
||
# 使用示例
|
||
if __name__ == "__main__":
|
||
client = ComfyUIClient(server_url="https://image.ai.faceta.cn")
|
||
|
||
input_image_path = "/Users/wandou/Downloads/aa.png"
|
||
|
||
with open('FaceImageArtView.json', 'r', encoding='utf-8') as f:
|
||
workflow = json.load(f)
|
||
|
||
# 提交任务
|
||
prompt_id = client.submit_prompt(workflow, input_image_path)
|
||
|
||
# 等待任务完成
|
||
output_images = client.wait_for_completion(prompt_id)
|
||
|
||
# 如果有生成的图片，获取并保存
|
||
if output_images:
|
||
for i, img_file in enumerate(output_images):
|
||
save_path = f"output_{i}.png"
|
||
client.get_image(img_file['filename'], save_path)
|