comfyui-discord/lib/comfyui.py

46 lines
1.6 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())
return output_images