From 6a4cf055399fe4625a3fde651b07791527197aba Mon Sep 17 00:00:00 2001 From: zeus Date: Fri, 3 May 2024 16:09:15 -0500 Subject: [PATCH] Changed get_images(ws, prompt) to asynchronous method to prevent lost packets --- script_examples/websockets_api_example.py | 49 ++++++++++++++--------- 1 file changed, 29 insertions(+), 20 deletions(-) diff --git a/script_examples/websockets_api_example.py b/script_examples/websockets_api_example.py index 57a6cbd9..958c5d12 100644 --- a/script_examples/websockets_api_example.py +++ b/script_examples/websockets_api_example.py @@ -6,6 +6,7 @@ import uuid import json import urllib.request import urllib.parse +import asyncio server_address = "127.0.0.1:8188" client_id = str(uuid.uuid4()) @@ -26,17 +27,23 @@ 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): +async def receive(ws): + return ws.recv() + +async def get_images(ws, prompt): prompt_id = queue_prompt(prompt)['prompt_id'] output_images = {} while True: - out = ws.recv() + out = await receive(ws) 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 + elif message['type'] == 'progress': + data = message['data'] + print("Processed step " + str(data['value']) + " out of " + str(data['max'])) else: continue #previews are binary data @@ -141,24 +148,26 @@ prompt_text = """ } } """ +async def api_example(): + prompt = json.loads(prompt_text) + #set the text prompt for our positive CLIPTextEncode + prompt["6"]["inputs"]["text"] = "masterpiece best quality man" + + #set the seed for our KSampler node + prompt["3"]["inputs"]["seed"] = 5 + + ws = websocket.WebSocket() + ws.connect("ws://{}/ws?clientId={}".format(server_address, client_id)) + images = await get_images(ws, prompt) -prompt = json.loads(prompt_text) -#set the text prompt for our positive CLIPTextEncode -prompt["6"]["inputs"]["text"] = "masterpiece best quality man" - -#set the seed for our KSampler node -prompt["3"]["inputs"]["seed"] = 5 - -ws = websocket.WebSocket() -ws.connect("ws://{}/ws?clientId={}".format(server_address, client_id)) -images = get_images(ws, prompt) - -#Commented out code to display the output images: + #Commented out code to display the output images: -# for node_id in images: -# for image_data in images[node_id]: -# from PIL import Image -# import io -# image = Image.open(io.BytesIO(image_data)) -# image.show() + # for node_id in images: + # for image_data in images[node_id]: + # from PIL import Image + # import io + # image = Image.open(io.BytesIO(image_data)) + # image.show() +loop=asyncio.get_event_loop() +loop.run_until_complete(api_example())