130 lines
5.6 KiB
Python
130 lines
5.6 KiB
Python
import json
|
|
import websocket
|
|
import uuid
|
|
import urllib.request
|
|
import urllib.parse
|
|
import random
|
|
import pandas as pd
|
|
|
|
COMFYUI_ENDPOINT = '127.0.0.1:8188'
|
|
|
|
class ComfyUITool:
|
|
def __init__(self, server_address, seed, workflowfile, working_dir):
|
|
self.server_address = server_address
|
|
self.client_id = str(uuid.uuid4()) # 生成一个唯一的客户端ID
|
|
self.seed = seed
|
|
self.workflowfile = workflowfile
|
|
self.working_dir = working_dir
|
|
|
|
def show_gif(self, fname):
|
|
import base64
|
|
from IPython import display
|
|
with open(fname, 'rb') as fd:
|
|
b64 = base64.b64encode(fd.read()).decode('ascii')
|
|
return display.HTML(f'<img src="data:image/gif;base64,{b64}" />')
|
|
|
|
def queue_prompt(self, prompt):
|
|
p = {"prompt": prompt, "client_id": self.client_id}
|
|
data = json.dumps(p).encode('utf-8')
|
|
req = urllib.request.Request("http://{}/prompt".format(self.server_address), data=data)
|
|
return json.loads(urllib.request.urlopen(req).read())
|
|
|
|
def get_image(self, filename, subfolder, folder_type):
|
|
data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
|
|
url_values = urllib.parse.urlencode(data)
|
|
with urllib.request.urlopen("http://{}/view?{}".format(self.server_address, url_values)) as response:
|
|
return response.read()
|
|
|
|
def get_history(self, prompt_id):
|
|
with urllib.request.urlopen("http://{}/history/{}".format(self.server_address, prompt_id)) as response:
|
|
return json.loads(response.read())
|
|
|
|
def get_images(self, ws, prompt):
|
|
prompt_id = self.queue_prompt(prompt)['prompt_id']
|
|
output_images = {}
|
|
while True:
|
|
out = ws.recv()
|
|
if isinstance(out, str):
|
|
message = json.loads(out)
|
|
if message['type'] == 'executing':
|
|
data = message['data']
|
|
if data['node'] is None and data['prompt_id'] == prompt_id:
|
|
print('执行完成')
|
|
break
|
|
else:
|
|
continue
|
|
|
|
history = self.get_history(prompt_id)[prompt_id]
|
|
for o in history['outputs']:
|
|
for node_id in history['outputs']:
|
|
node_output = history['outputs'][node_id]
|
|
# 图片分支
|
|
if 'images' in node_output:
|
|
images_output = []
|
|
for image in node_output['images']:
|
|
image_data = self.get_image(image['filename'], image['subfolder'], image['type'])
|
|
images_output.append(image_data)
|
|
output_images[node_id] = images_output
|
|
# 视频分支
|
|
if 'videos' in node_output:
|
|
videos_output = []
|
|
for video in node_output['videos']:
|
|
video_data = self.get_image(video['filename'], video['subfolder'], video['type'])
|
|
videos_output.append(video_data)
|
|
output_images[node_id] = videos_output
|
|
print('获取图片完成')
|
|
return output_images
|
|
|
|
def parse_worflow(self, ws, prompt):
|
|
# 获取工作流文件路径
|
|
workflowfile = self.workflowfile
|
|
print('workflowfile:' + workflowfile)
|
|
# 打开工作流文件并加载JSON数据
|
|
with open(workflowfile, 'r', encoding="utf-8") as workflow_api_txt2gif_file:
|
|
prompt_data = json.load(workflow_api_txt2gif_file)
|
|
### 注意这里根据自己的工作流修改, 这里是替换prompt 节点的代码。
|
|
prompt_data["6"]["inputs"]["text"] = prompt
|
|
return self.get_images(ws, prompt_data)
|
|
|
|
def generate_clip(self, prompt, idx=1):
|
|
ws = websocket.WebSocket()
|
|
ws.connect("ws://{}/ws?clientId={}".format(self.server_address, self.client_id))
|
|
images = self.parse_worflow(ws, prompt)
|
|
for node_id in images:
|
|
for image_data in images[node_id]:
|
|
from datetime import datetime
|
|
# 获取当前时间，并格式化为 YYYYMMDDHHMMSS 的格式
|
|
timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
|
|
# 使用格式化的时间戳在文件名中
|
|
GIF_LOCATION = "{}/{}_{}_{}.png".format(self.working_dir, idx, self.seed, timestamp)
|
|
with open(GIF_LOCATION, "wb") as binary_file:
|
|
# 写入二进制文件
|
|
binary_file.write(image_data)
|
|
self.show_gif(GIF_LOCATION)
|
|
def read_prompts_from_excel(self, csv_file_path):
|
|
df = pd.read_excel(csv_file_path)
|
|
return df['prompt'].tolist()
|
|
|
|
server_address = '127.1.1.1:8188'
|
|
workflow_file = './workflows/flux_redux.json'
|
|
workfolw_seed = 162434675638754 # workflowfile 开头总的中的seed
|
|
output_dir= 'output'
|
|
|
|
# 创建 ComfyUITool 实例
|
|
comfyui_tool = ComfyUITool(server_address, workfolw_seed, workflow_file, output_dir)
|
|
prompt = "sunset,The night had not fully lifted,and the coastal city was still slumbering. \
|
|
The waves gently lapped the beach,playing a tender overture. A faint light emerged \
|
|
on the horizon,gradually changing from pale white to golden yellow,slowly outlining \
|
|
the city's silhouette. "
|
|
|
|
## 更换提示词生成图片
|
|
comfyui_tool.generate_clip(prompt)
|
|
|
|
# 读取excel中多条提示词(每条提示词一行, 第一行为"prompt") 批量生成图片
|
|
prompts_file_path = 'prompt.xlsx'
|
|
prompts = comfyui_tool.read_prompts_from_excel(prompts_file_path)
|
|
print(prompts)
|
|
idx = 1
|
|
for prompt in prompts:
|
|
comfyui_tool.generate_clip(prompt, idx)
|
|
idx += 1 |