first steps ws api

Author: phDoll

.gitignore

@@ -0,0 +1,2 @@
+output*
+__pycache__/

api/__init__.py

@@ -0,0 +1,86 @@
+#This is an example that uses the websockets api to know when a prompt execution is done
+#Once the prompt execution is done it downloads the images using the /history endpoint
+
+import uuid
+import json
+import urllib.request
+import urllib.parse
+import websocket #NOTE: websocket-client (https://github.com/websocket-client/websocket-client)
+from PIL import Image
+import io
+
+
+server_address='127.0.0.1:8188'
+client_id=str(uuid.uuid4())
+
+ws = websocket.WebSocket()
+ws.connect("ws://{}/ws?clientId={}".format(server_address, client_id))
+
+def generate_image_by_prompt(prompt, output_path, save_previews=False):
+    prompt_id = queue_prompt(prompt)['prompt_id']
+    images = get_images(ws, prompt_id, save_previews)
+    save_image(images, output_path, save_previews)
+
+def save_image(images, output_path, save_previews):
+    for itm in images:
+        if itm['type'] == 'temp' and save_previews:
+            image = Image.open(io.BytesIO(itm['image_data']))
+            image.save(output_path + 'temp/' + itm['file_name'])
+        else:
+            image = Image.open(io.BytesIO(itm['image_data']))
+            image.save(output_path + itm['file_name'])
+
+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)
+    return json.loads(urllib.request.urlopen(req).read())
+
+def interupt_prompt():
+    req =  urllib.request.Request("http://{}/interrupt".format(server_address), data={})
+    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_id, allow_preview = False):
+    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 node_id in history['outputs']:
+        node_output = history['outputs'][node_id]
+        output_data = {}
+        if 'images' in node_output:
+            for image in node_output['images']:
+                if allow_preview and image['type'] == 'temp':
+                    preview_data = get_image(image['filename'], image['subfolder'], image['type'])
+                    output_data['image_data'] = preview_data
+                if image['type'] == 'output':
+                    image_data = get_image(image['filename'], image['subfolder'], image['type'])
+                    output_data['image_data'] = image_data
+        output_data['file_name'] = image['filename']
+        output_data['type'] = image['type']
+        output_images.append(output_data)
+
+    return output_images
+
+
+

api/websocket_api.py

@@ -0,0 +1,22 @@
+from utils.actions.prompt_to_image import prompt_to_image
+from utils.actions.load_workflow import load_workflow
+from utils.actions.interrupt_prompt import interrupt
+import time
+import sys
+
+def main():
+    try:
+      print("Welcome to the program!")
+      workflow = load_workflow('./workflows/test_workflow.json')
+
+      prompt_to_image(workflow, 'a tiny little green and blue monster wearing a coat and a little hat, very cute with sharp teeths', True)
+    except Exception as e:
+      print(f"An error occurred: {e}")
+      exit_program()
+
+def exit_program():
+  print("Exiting the program...")
+  interrupt()
+  sys.exit(0)
+
+main()

main.py

@@ -0,0 +1,3 @@
+Pillow==10.0.0
+Pillow==10.2.0
+websocket_client==1.7.0

requirements.txt

@@ -0,0 +1,3 @@
+from api.websocket_api import interupt_prompt
+def interrupt():
+  interupt_prompt()

utils/__init__.py

@@ -0,0 +1,8 @@
+import json
+
+def load_workflow(workflow_path):
+  file = open(workflow_path)
+  workflow = json.load(file)
+  workflow = json.dumps(workflow)
+
+  return workflow

utils/actions/__init__.py

@@ -0,0 +1,23 @@
+from api.websocket_api import generate_image_by_prompt
+from utils.helpers.find_node import find_node
+from utils.helpers.randomize_seed import generate_random_15_digit_number
+from utils.helpers.replace_key import replace_key
+import json
+def prompt_to_image(workflow, positve_prompt, negative_prompt='', save_previews=False):
+  prompt = json.loads(workflow)
+  replace_key(prompt, 'seed', generate_random_15_digit_number())
+  postive_prompt_id = find_node(prompt, 'positive')[0]
+  positive_prompt_node = find_node(prompt, postive_prompt_id)
+
+  if negative_prompt != '':
+    negative_prompt_id = find_node(prompt, 'negative')[0]
+    negative_prompt_node = find_node(prompt, negative_prompt_id)
+    negative_prompt_node['inputs']['text'] = negative_prompt
+
+  positive_prompt_node['inputs']['text'] = positve_prompt
+
+  generate_image_by_prompt(prompt, './output/', save_previews)
+
+
+
+