setup api with multiple classes

Author: phDoll

api/api_helpers.py

@@ -0,0 +1,95 @@
+import json
+from PIL import Image
+import io
+import os
+
+# Assuming the import paths are correct and the methods are defined elsewhere:
+from api.websocket_api import queue_prompt, get_history, get_image, upload_image, clear_comfy_cache
+from api.open_websocket import open_websocket_connection
+
+def generate_image_by_prompt(prompt, output_path, save_previews=False):
+  try:
+    ws, server_address, client_id = open_websocket_connection()
+    prompt_id = queue_prompt(prompt, client_id, server_address)['prompt_id']
+    track_progress(prompt, ws, prompt_id)
+    images = get_images(prompt_id, server_address, save_previews)
+    save_image(images, output_path, save_previews)
+  finally:
+    ws.close()
+
+def generate_image_by_prompt_and_image(prompt, output_path, input_path, filename, save_previews=False):
+  try:
+    ws, server_address, client_id = open_websocket_connection()
+    upload_image(input_path, filename, server_address, client_id)
+    prompt_id = queue_prompt(prompt, client_id, server_address)['prompt_id']
+    track_progress(prompt, ws, prompt_id)
+    images = get_images(prompt_id, server_address, save_previews)
+    save_image(images, output_path, save_previews)
+  finally:
+    ws.close()
+
+def save_image(images, output_path, save_previews):
+    for itm in images:
+        directory = os.path.join(output_path, 'temp/') if itm['type'] == 'temp' and save_previews else output_path
+        os.makedirs(directory, exist_ok=True)
+        try:
+            image = Image.open(io.BytesIO(itm['image_data']))
+            image.save(os.path.join(directory, itm['file_name']))
+        except Exception as e:
+            print(f"Failed to save image {itm['file_name']}: {e}")
+
+def track_progress(prompt, ws, prompt_id):
+  node_ids = list(prompt.keys())
+  finished_nodes = []
+
+  while True:
+      out = ws.recv()
+      if isinstance(out, str):
+          message = json.loads(out)
+          if message['type'] == 'progress':
+              data = message['data']
+              current_step = data['value']
+              print('In K-Sampler -> Step: ', current_step, ' of: ', data['max'])
+          if message['type'] == 'execution_cached':
+              data = message['data']
+              for itm in data['nodes']:
+                  if itm not in finished_nodes:
+                      finished_nodes.append(itm)
+                      print('Progess: ', len(finished_nodes), '/', len(node_ids), ' Tasks done')
+          if message['type'] == 'executing':
+              data = message['data']
+              if data['node'] not in finished_nodes:
+                  finished_nodes.append(data['node'])
+                  print('Progess: ', len(finished_nodes), '/', len(node_ids), ' Tasks done')
+
+
+              if data['node'] is None and data['prompt_id'] == prompt_id:
+                  break #Execution is done
+      else:
+          continue #previews are binary data
+  return
+
+def get_images(prompt_id, server_address, allow_preview = False):
+  output_images = []
+
+  history = get_history(prompt_id, server_address)[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'], server_address)
+                  output_data['image_data'] = preview_data
+              if image['type'] == 'output':
+                  image_data = get_image(image['filename'], image['subfolder'], image['type'], server_address)
+                  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
+
+
+def clear():
+  clear_comfy_cache()

api/open_websocket.py

@@ -0,0 +1,10 @@
+import websocket #NOTE: websocket-client (https://github.com/websocket-client/websocket-client)
+import uuid
+
+def open_websocket_connection():
+  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))
+  return ws, server_address, client_id

api/websocket_api.py

@@ -1,137 +1,56 @@
-#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
-from utils.helpers.find_node import find_node
 from requests_toolbelt import MultipartEncoder
 
-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 setup_image(input_path, name, type="input", overwrite=False):
-    with open(input_path, 'rb') as file:
-        multipart_data = MultipartEncoder(
-            fields={
-                'image': (name, file, 'image/jpeg'),  # Adjust the content-type accordingly
-                'type': type,
-                'overwrite': str(overwrite).lower()
-            }
-        )
-
-        data = multipart_data
-        headers = {'Content-Type': multipart_data.content_type}
-        request = urllib.request.Request("http://{}/upload/image".format(server_address), data=data, headers=headers)
-        with urllib.request.urlopen(request) as response:
-            return response.read()
-
-def generate_image_by_prompt(prompt, output_path, save_previews=False):
-    prompt_id = queue_prompt(prompt)['prompt_id']
-    track_progress(prompt, ws, 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 track_progress(prompt, ws, prompt_id):
-    node_ids = list(prompt.keys())
-    finished_nodes = []
-
-    while True:
-        out = ws.recv()
-        if isinstance(out, str):
-            message = json.loads(out)
-            if message['type'] == 'progress':
-                data = message['data']
-                current_step = data['value']
-                print('In K-Sampler -> Step: ', current_step, ' of: ', data['max'])
-            if message['type'] == 'execution_cached':
-                data = message['data']
-                for itm in data['nodes']:
-                    if itm not in finished_nodes:
-                        finished_nodes.append(itm)
-                        print('Progess: ', len(finished_nodes), '/', len(node_ids), ' Tasks done')
-            if message['type'] == 'executing':
-                data = message['data']
-                if data['node'] not in finished_nodes:
-                    finished_nodes.append(data['node'])
-                    print('Progess: ', len(finished_nodes), '/', len(node_ids), ' Tasks done')
-
-
-                if data['node'] is None and data['prompt_id'] == prompt_id:
-                    break #Execution is done
-        else:
-            continue #previews are binary data
-    return
-
-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 = []
-
-    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
-
-def get_node_info_by_class(node_class):
-    with urllib.request.urlopen("http://{}/object_info/{}".format(server_address, node_class)) as response:
-        return json.loads(response.read())
-
-def clear_comfy_cache(unload_models=False, free_memory=False):
-    clear_data = {
-        "unload_models": unload_models,
-        "free_memory": free_memory
-    }
-    data = json.dumps(clear_data).encode('utf-8')
-
-    with urllib.request.urlopen("http://{}/free".format(server_address), data=data) as response:
-        return response.read()
+def upload_image(input_path, name, server_address, image_type="input", overwrite=False):
+	with open(input_path, 'rb') as file:
+		multipart_data = MultipartEncoder(
+      fields= {
+        'image': (name, file, 'image/png'),
+        'type': image_type,
+        'overwrite': str(overwrite).lower()
+      }
+		)
+
+		data = multipart_data
+		headers = { 'Content-Type': multipart_data.content_type }
+		request = urllib.request.Request("http://{}/upload/image".format(server_address), data=data, headers=headers)
+		with urllib.request.urlopen(request) as response:
+			return response.read()
+
+def queue_prompt(prompt, client_id, server_address):
+  p = {"prompt": prompt, "client_id": client_id}
+  headers = {'Content-Type': 'application/json'}
+  data = json.dumps(p).encode('utf-8')
+  req =  urllib.request.Request("http://{}/prompt".format(server_address), data=data, headers=headers)
+  return json.loads(urllib.request.urlopen(req).read())
+
+def interupt_prompt(server_address):
+  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, server_address):
+  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, server_address):
+  with urllib.request.urlopen("http://{}/history/{}".format(server_address, prompt_id)) as response:
+      return json.loads(response.read())
+
+def get_node_info_by_class(node_class, server_address):
+  with urllib.request.urlopen("http://{}/object_info/{}".format(server_address, node_class)) as response:
+      return json.loads(response.read())
+
+def clear_comfy_cache(server_address, unload_models=False, free_memory=False):
+  clear_data = {
+    "unload_models": unload_models,
+    "free_memory": free_memory
+  }
+  data = json.dumps(clear_data).encode('utf-8')
+
+  with urllib.request.urlopen("http://{}/free".format(server_address), data=data) as response:
+    return response.read()
 

main.py

@@ -1,19 +1,18 @@
 from utils.actions.prompt_to_image import prompt_to_image
 from utils.actions.prompt_image_to_image import prompt_image_to_image
 from utils.actions.load_workflow import load_workflow
-from api.websocket_api import clear_comfy_cache
-from api.websocket_api import get_image
-import time
+from api.api_helpers import clear
 import sys
 
 def main():
     try:
       print("Welcome to the program!")
-      workflow = load_workflow('./workflows/image_to_image.json')
-
-      # prompt_to_image(workflow, 'beautiful woman sitting on a desk in a nice restaurant, candlelight dinner atmosphere, wearing a red dress', save_previews=True)
-      input_path = './input/ComfyUI_00103_.png'
-      prompt_image_to_image(workflow, input_path, 'beautiful [white woman], (dark lighting), curly blond hair', save_previews=True)
+      workflow = load_workflow('./workflows/base_workflow.json')
+      for iter in range(1, 11):
+        prompt_to_image(workflow, '(realistic:1.25), beautiful:1.1) mountain landscape with a deep blue lake, photolike, high detail, monoton colors', 'lowres, text, branding, watermark, humans, frames, painting', save_previews=True)
+      # prompt_to_image(workflow, '(beautiful woman:1.3) sitting on a desk in a nice restaurant with a (glass of wine and plate with salat:0.9), (candlelight dinner atmosphere:1.1), (wearing a red evening dress:1.2), dimmed lighting, cinema, high detail', save_previews=True)
+      # input_path = './input/ComfyUI_00241_.png'
+      # prompt_image_to_image(workflow, input_path, '(white woman wearing a black evening dress:1.5), dimmed lighting, cinema, high detail', save_previews=True)
     except Exception as e:
       print(f"An error occurred: {e}")
       exit_program()
@@ -22,13 +21,7 @@ def exit_program():
   print("Exiting the program...")
   sys.exit(0)
 
-def clear():
-  clear_comfy_cache(True, True)
-
-def image_by_file():
-  print(get_image('ComfyUI_00042_.png', '', 'output'))
+def clear_comfy():
+  clear(True, True)
 
 main()
-# clear()
-
-# image_by_file()

utils/actions/load_workflow.py

@@ -1,8 +1,13 @@
 import json
 
 def load_workflow(workflow_path):
-  file = open(workflow_path)
-  workflow = json.load(file)
-  workflow = json.dumps(workflow)
-
-  return workflow
+    try:
+        with open(workflow_path, 'r') as file:
+            workflow = json.load(file)
+            return json.dumps(workflow)
+    except FileNotFoundError:
+        print(f"The file {workflow_path} was not found.")
+        return None
+    except json.JSONDecodeError:
+        print(f"The file {workflow_path} contains invalid JSON.")
+        return None

utils/actions/prompt_image_to_image.py

@@ -1,16 +1,11 @@
-from api.websocket_api import generate_image_by_prompt
-from api.websocket_api import setup_image
-from utils.helpers.find_node import find_node
+from api.api_helpers import generate_image_by_prompt_and_image
 from utils.helpers.randomize_seed import generate_random_15_digit_number
-from utils.helpers.replace_key import replace_key
-from utils.helpers.find_parent import find_parent_of_key
 import json
 def prompt_image_to_image(workflow, input_path, positve_prompt, negative_prompt='', save_previews=False):
   prompt = json.loads(workflow)
-  replace_key(prompt, 'seed', generate_random_15_digit_number())
   id_to_class_type = {id: details['class_type'] for id, details in prompt.items()}
   k_sampler = [key for key, value in id_to_class_type.items() if value == 'KSampler'][0]
-
+  prompt.get(k_sampler)['inputs']['seed'] = generate_random_15_digit_number()
   postive_input_id = prompt.get(k_sampler)['inputs']['positive'][0]
   prompt.get(postive_input_id)['inputs']['text_g'] = positve_prompt
   prompt.get(postive_input_id)['inputs']['text_l'] = positve_prompt
@@ -20,14 +15,9 @@ def prompt_image_to_image(workflow, input_path, positve_prompt, negative_prompt=
     id_to_class_type.get(negative_input_id)['inputs']['text_g'] = negative_prompt
     id_to_class_type.get(negative_input_id)['inputs']['text_l'] = negative_prompt
 
-
   image_loader = [key for key, value in id_to_class_type.items() if value == 'LoadImage'][0]
   filename = input_path.split('/')[-1]
   prompt.get(image_loader)['inputs']['image'] = filename
 
-  setup_image(input_path, filename)
-  generate_image_by_prompt(prompt, './output/', save_previews)
-
-
-
+  generate_image_by_prompt_and_image(prompt, './output/', input_path, filename, save_previews)