62 lines
2.1 KiB
Python
62 lines
2.1 KiB
Python
import time
|
|
|
|
import asyncio
|
|
import aiohttp
|
|
|
|
from urllib.request import urlopen as urlopen
|
|
from urllib.parse import urlencode as urlencode
|
|
|
|
async def get_comfyui_generations(api_url, workflow):
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.post(f"{api_url}/prompt", json={"prompt": workflow}) as resp:
|
|
resp_json = await resp.json()
|
|
|
|
prompt_id = resp_json["prompt_id"]
|
|
|
|
# Loop endlessly until ComfyUI confirms our prompt's completion:
|
|
while True:
|
|
async with session.get(f"{api_url}/history/{prompt_id}") as resp:
|
|
resp_json = await resp.json()
|
|
if not resp_json:
|
|
time.sleep(1)
|
|
continue
|
|
break
|
|
|
|
# Read the output history anmd fetch each image:
|
|
history = resp_json[prompt_id]
|
|
output_images = []
|
|
|
|
for o in history["outputs"]:
|
|
for node_id in history["outputs"]:
|
|
node_output = history["outputs"][node_id]
|
|
files_output = []
|
|
|
|
if "images" in node_output:
|
|
for image in node_output["images"]:
|
|
url_params = urlencode({
|
|
"filename": image["filename"],
|
|
"subfolder": image["subfolder"],
|
|
"type": image["type"]
|
|
})
|
|
|
|
async with session.get(f"{api_url}/view?{url_params}") as resp:
|
|
output_images.append(await resp.content.read())
|
|
|
|
|
|
messages = resp_json[prompt_id]["status"]["messages"]
|
|
timestamp_start = 0
|
|
timestamp_finish = 0
|
|
|
|
for m in messages:
|
|
label = m[0]
|
|
obj = m[1]
|
|
if label == "execution_start": timestamp_start = obj["timestamp"]
|
|
if label == "execution_success": timestamp_finish = obj["timestamp"]
|
|
|
|
execution_time = (timestamp_finish - timestamp_start) / 1000.0
|
|
|
|
return {
|
|
"execution_time": execution_time,
|
|
"images": output_images
|
|
}
|