Skip to content

Instantly share code, notes, and snippets.

@trojblue
Created June 30, 2024 19:37
Show Gist options
  • Save trojblue/acd6f026ef3669af6a81d107e259749a to your computer and use it in GitHub Desktop.
Save trojblue/acd6f026ef3669af6a81d107e259749a to your computer and use it in GitHub Desktop.
  1. enable "dev usage" in comfy settings to export api_workflow.json:
  1. use the script:
import websocket #NOTE: websocket-client (https://github.com/websocket-client/websocket-client)
import uuid
import json
import urllib.request
import urllib.parse

from PIL import Image
import io

server_address = "127.0.0.1:8188"
client_id = str(uuid.uuid4())

def queue_prompt(prompt):
    p = {"prompt": prompt, "client_id": client_id}
    data = json.dumps(p).encode('utf-8')
    req =  urllib.request.Request("http://{}/prompt".format(server_address), data=data,method="POST")
    return json.loads(urllib.request.urlopen(req).read())

def get_image(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(server_address, url_values)) as response:
        return response.read()

def get_history(prompt_id):
    with urllib.request.urlopen("http://{}/history/{}".format(server_address, prompt_id)) as response:
        return json.loads(response.read())

def get_images(ws, prompt):
    prompt_id = 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:
                    break #Execution is done
        else:
            continue #previews are binary data

    history = 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 = get_image(image['filename'], image['subfolder'], image['type'])
                    images_output.append(image_data)
            output_images[node_id] = images_output

    return output_images

with open("../workflow_api.json","r",encoding="utf-8") as f:
    workflow_jsondata = f.read()

jsonwf=json.loads(workflow_jsondata)


ws = websocket.WebSocket()
ws.connect("ws://{}/ws?clientId={}".format(server_address, client_id))
images = get_images(ws, jsonwf)


for node_id in images:
    i=0
    for image_data in images[node_id]:
        image = Image.open(io.BytesIO(image_data))
        image.save(f"output-{i}-{node_id}.webp", "WEBP")
        i+=1
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment