File size: 7,777 Bytes
4450790 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 |
import os
from typing import Sequence, Mapping, Any, Union
import sys
sys.path.append('../')
args = None
def import_custom_nodes() -> None:
"""Find all custom nodes in the custom_nodes folder and add those node objects to NODE_CLASS_MAPPINGS
This function sets up a new asyncio event loop, initializes the PromptServer,
creates a PromptQueue, and initializes the custom nodes.
"""
import asyncio
import execution
from nodes import init_extra_nodes
import server
# Creating a new event loop and setting it as the default loop
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
# Creating an instance of PromptServer with the loop
server_instance = server.PromptServer(loop)
execution.PromptQueue(server_instance)
# Initializing custom nodes
init_extra_nodes(init_custom_nodes=True)
def find_path(name: str, path: str = None) -> str:
"""
Recursively looks at parent folders starting from the given path until it finds the given name.
Returns the path as a Path object if found, or None otherwise.
"""
# If no path is given, use the current working directory
if path is None:
if args is None or args.comfyui_directory is None:
path = os.getcwd()
else:
path = args.comfyui_directory
# Check if the current directory contains the name
if name in os.listdir(path):
path_name = os.path.join(path, name)
print(f"{name} found: {path_name}")
return path_name
# Get the parent directory
parent_directory = os.path.dirname(path)
# If the parent directory is the same as the current directory, we've reached the root and stop the search
if parent_directory == path:
return None
# Recursively call the function with the parent directory
return find_path(name, parent_directory)
def add_comfyui_directory_to_sys_path() -> None:
"""
Add 'ComfyUI' to the sys.path
"""
comfyui_path = find_path('ComfyUI')
if comfyui_path is not None and os.path.isdir(comfyui_path):
sys.path.append(comfyui_path)
import __main__
if getattr(__main__, "__file__", None) is None:
__main__.__file__ = os.path.join(comfyui_path, "main.py")
print(f"'{comfyui_path}' added to sys.path")
def add_extra_model_paths() -> None:
"""
Parse the optional extra_model_paths.yaml file and add the parsed paths to the sys.path.
"""
from utils.extra_config import load_extra_path_config
extra_model_paths = find_path("extra_model_paths.yaml")
if extra_model_paths is not None:
load_extra_path_config(extra_model_paths)
else:
print("Could not find the extra_model_paths config file.")
def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
"""Returns the value at the given index of a sequence or mapping.
If the object is a sequence (like list or string), returns the value at the given index.
If the object is a mapping (like a dictionary), returns the value at the index-th key.
Some return a dictionary, in these cases, we look for the "results" key
Args:
obj (Union[Sequence, Mapping]): The object to retrieve the value from.
index (int): The index of the value to retrieve.
Returns:
Any: The value at the given index.
Raises:
IndexError: If the index is out of bounds for the object and the object is not a mapping.
"""
try:
return obj[index]
except KeyError:
return obj['result'][index]
def parse_arg(s: Any):
""" Parses a JSON string, returning it unchanged if the parsing fails. """
if __name__ == "__main__" or not isinstance(s, str):
return s
try:
return json.loads(s)
except json.JSONDecodeError:
return s
def save_image_wrapper(context, cls):
if args.output is None: return cls
from PIL import Image, ImageOps, ImageSequence
from PIL.PngImagePlugin import PngInfo
import numpy as np
class WrappedSaveImage(cls):
counter = 0
def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
if args.output is None:
return super().save_images(images, filename_prefix, prompt, extra_pnginfo)
else:
if len(images) > 1 and args.output == "-":
raise ValueError("Cannot save multiple images to stdout")
filename_prefix += self.prefix_append
results = list()
for (batch_number, image) in enumerate(images):
i = 255. * image.cpu().numpy()
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
metadata = None
if not args.disable_metadata:
metadata = PngInfo()
if prompt is not None:
metadata.add_text("prompt", json.dumps(prompt))
if extra_pnginfo is not None:
for x in extra_pnginfo:
metadata.add_text(x, json.dumps(extra_pnginfo[x]))
if args.output == "-":
# Hack to briefly restore stdout
if context is not None:
context.__exit__(None, None, None)
try:
img.save(sys.stdout.buffer, format="png", pnginfo=metadata, compress_level=self.compress_level)
finally:
if context is not None:
context.__enter__()
else:
subfolder = ""
if len(images) == 1:
if os.path.isdir(args.output):
subfolder = args.output
file = "output.png"
else:
subfolder, file = os.path.split(args.output)
if subfolder == "":
subfolder = os.getcwd()
else:
if os.path.isdir(args.output):
subfolder = args.output
file = filename_prefix
else:
subfolder, file = os.path.split(args.output)
if subfolder == "":
subfolder = os.getcwd()
files = os.listdir(subfolder)
file_pattern = file
while True:
filename_with_batch_num = file_pattern.replace("%batch_num%", str(batch_number))
file = f"{filename_with_batch_num}_{self.counter:05}.png"
self.counter += 1
if file not in files:
break
img.save(os.path.join(subfolder, file), pnginfo=metadata, compress_level=self.compress_level)
print("Saved image to", os.path.join(subfolder, file))
results.append({
"filename": file,
"subfolder": subfolder,
"type": self.type
})
return {"ui": {"images": results}}
return WrappedSaveImage
|