brunvelop/ComfyUI

1#This is an example that uses the websockets api to know when a prompt execution is done2#Once the prompt execution is done it downloads the images using the /history endpoint3 4import websocket #NOTE: websocket-client (https://github.com/websocket-client/websocket-client)5import uuid6import json7import urllib.request8import urllib.parse9 10server_address = "127.0.0.1:8188"11client_id = str(uuid.uuid4())12 13def queue_prompt(prompt):14    p = {"prompt": prompt, "client_id": client_id}15    data = json.dumps(p).encode('utf-8')16    req =  urllib.request.Request("http://{}/prompt".format(server_address), data=data)17    return json.loads(urllib.request.urlopen(req).read())18 19def get_image(filename, subfolder, folder_type):20    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}21    url_values = urllib.parse.urlencode(data)22    with urllib.request.urlopen("http://{}/view?{}".format(server_address, url_values)) as response:23        return response.read()24 25def get_history(prompt_id):26    with urllib.request.urlopen("http://{}/history/{}".format(server_address, prompt_id)) as response:27        return json.loads(response.read())28 29def get_images(ws, prompt):30    prompt_id = queue_prompt(prompt)['prompt_id']31    output_images = {}32    while True:33        out = ws.recv()34        if isinstance(out, str):35            message = json.loads(out)36            if message['type'] == 'executing':37                data = message['data']38                if data['node'] is None and data['prompt_id'] == prompt_id:39                    break #Execution is done40        else:41            continue #previews are binary data42 43    history = get_history(prompt_id)[prompt_id]44    for o in history['outputs']:45        for node_id in history['outputs']:46            node_output = history['outputs'][node_id]47            if 'images' in node_output:48                images_output = []49                for image in node_output['images']:50                    image_data = get_image(image['filename'], image['subfolder'], image['type'])51                    images_output.append(image_data)52            output_images[node_id] = images_output53 54    return output_images55 56prompt_text = """57{58    "3": {59        "class_type": "KSampler",60        "inputs": {61            "cfg": 8,62            "denoise": 1,63            "latent_image": [64                "5",65                066            ],67            "model": [68                "4",69                070            ],71            "negative": [72                "7",73                074            ],75            "positive": [76                "6",77                078            ],79            "sampler_name": "euler",80            "scheduler": "normal",81            "seed": 8566257,82            "steps": 2083        }84    },85    "4": {86        "class_type": "CheckpointLoaderSimple",87        "inputs": {88            "ckpt_name": "v1-5-pruned-emaonly.ckpt"89        }90    },91    "5": {92        "class_type": "EmptyLatentImage",93        "inputs": {94            "batch_size": 1,95            "height": 512,96            "width": 51297        }98    },99    "6": {100        "class_type": "CLIPTextEncode",101        "inputs": {102            "clip": [103                "4",104                1105            ],106            "text": "masterpiece best quality girl"107        }108    },109    "7": {110        "class_type": "CLIPTextEncode",111        "inputs": {112            "clip": [113                "4",114                1115            ],116            "text": "bad hands"117        }118    },119    "8": {120        "class_type": "VAEDecode",121        "inputs": {122            "samples": [123                "3",124                0125            ],126            "vae": [127                "4",128                2129            ]130        }131    },132    "9": {133        "class_type": "SaveImage",134        "inputs": {135            "filename_prefix": "ComfyUI",136            "images": [137                "8",138                0139            ]140        }141    }142}143"""144 145prompt = json.loads(prompt_text)146#set the text prompt for our positive CLIPTextEncode147prompt["6"]["inputs"]["text"] = "masterpiece best quality man"148 149#set the seed for our KSampler node150prompt["3"]["inputs"]["seed"] = 5151 152ws = websocket.WebSocket()153ws.connect("ws://{}/ws?clientId={}".format(server_address, client_id))154images = get_images(ws, prompt)155 156#Commented out code to display the output images:157 158# for node_id in images:159#     for image_data in images[node_id]:160#         from PIL import Image161#         import io162#         image = Image.open(io.BytesIO(image_data))163#         image.show()164 165