|
import asyncio |
|
import os |
|
import json |
|
import shutil |
|
import inspect |
|
import aiohttp |
|
from server import PromptServer |
|
from tqdm import tqdm |
|
|
|
config = None |
|
|
|
|
|
def is_logging_enabled(): |
|
config = get_extension_config() |
|
if "logging" not in config: |
|
return False |
|
return config["logging"] |
|
|
|
|
|
def log(message, type=None, always=False, name=None): |
|
if not always and not is_logging_enabled(): |
|
return |
|
|
|
if type is not None: |
|
message = f"[{type}] {message}" |
|
|
|
if name is None: |
|
name = get_extension_config()["name"] |
|
|
|
print(f"(pysssss:{name}) {message}") |
|
|
|
|
|
def get_ext_dir(subpath=None, mkdir=False): |
|
dir = os.path.dirname(__file__) |
|
if subpath is not None: |
|
dir = os.path.join(dir, subpath) |
|
|
|
dir = os.path.abspath(dir) |
|
|
|
if mkdir and not os.path.exists(dir): |
|
os.makedirs(dir) |
|
return dir |
|
|
|
|
|
def get_comfy_dir(subpath=None, mkdir=False): |
|
dir = os.path.dirname(inspect.getfile(PromptServer)) |
|
if subpath is not None: |
|
dir = os.path.join(dir, subpath) |
|
|
|
dir = os.path.abspath(dir) |
|
|
|
if mkdir and not os.path.exists(dir): |
|
os.makedirs(dir) |
|
return dir |
|
|
|
|
|
def get_web_ext_dir(): |
|
config = get_extension_config() |
|
name = config["name"] |
|
dir = get_comfy_dir("web/extensions/pysssss") |
|
if not os.path.exists(dir): |
|
os.makedirs(dir) |
|
dir = os.path.join(dir, name) |
|
return dir |
|
|
|
|
|
def get_extension_config(reload=False): |
|
global config |
|
if reload == False and config is not None: |
|
return config |
|
|
|
config_path = get_ext_dir("pysssss.json") |
|
default_config_path = get_ext_dir("pysssss.default.json") |
|
if not os.path.exists(config_path): |
|
if os.path.exists(default_config_path): |
|
shutil.copy(default_config_path, config_path) |
|
if not os.path.exists(config_path): |
|
log(f"Failed to create config at {config_path}", type="ERROR", always=True, name="???") |
|
print(f"Extension path: {get_ext_dir()}") |
|
return {"name": "Unknown", "version": -1} |
|
|
|
else: |
|
log("Missing pysssss.default.json, this extension may not work correctly. Please reinstall the extension.", |
|
type="ERROR", always=True, name="???") |
|
print(f"Extension path: {get_ext_dir()}") |
|
return {"name": "Unknown", "version": -1} |
|
|
|
with open(config_path, "r") as f: |
|
config = json.loads(f.read()) |
|
return config |
|
|
|
|
|
def link_js(src, dst): |
|
src = os.path.abspath(src) |
|
dst = os.path.abspath(dst) |
|
if os.name == "nt": |
|
try: |
|
import _winapi |
|
_winapi.CreateJunction(src, dst) |
|
return True |
|
except: |
|
pass |
|
try: |
|
os.symlink(src, dst) |
|
return True |
|
except: |
|
import logging |
|
logging.exception('') |
|
return False |
|
|
|
|
|
def is_junction(path): |
|
if os.name != "nt": |
|
return False |
|
try: |
|
return bool(os.readlink(path)) |
|
except OSError: |
|
return False |
|
|
|
|
|
def install_js(): |
|
src_dir = get_ext_dir("web/js") |
|
if not os.path.exists(src_dir): |
|
log("No JS") |
|
return |
|
|
|
should_install = should_install_js() |
|
if should_install: |
|
log("it looks like you're running an old version of ComfyUI that requires manual setup of web files, it is recommended you update your installation.", "warning", True) |
|
dst_dir = get_web_ext_dir() |
|
linked = os.path.islink(dst_dir) or is_junction(dst_dir) |
|
if linked or os.path.exists(dst_dir): |
|
if linked: |
|
if should_install: |
|
log("JS already linked") |
|
else: |
|
os.unlink(dst_dir) |
|
log("JS unlinked, PromptServer will serve extension") |
|
elif not should_install: |
|
shutil.rmtree(dst_dir) |
|
log("JS deleted, PromptServer will serve extension") |
|
return |
|
|
|
if not should_install: |
|
log("JS skipped, PromptServer will serve extension") |
|
return |
|
|
|
if link_js(src_dir, dst_dir): |
|
log("JS linked") |
|
return |
|
|
|
log("Copying JS files") |
|
shutil.copytree(src_dir, dst_dir, dirs_exist_ok=True) |
|
|
|
|
|
def should_install_js(): |
|
return not hasattr(PromptServer.instance, "supports") or "custom_nodes_from_web" not in PromptServer.instance.supports |
|
|
|
|
|
def init(check_imports=None): |
|
log("Init") |
|
|
|
if check_imports is not None: |
|
import importlib.util |
|
for imp in check_imports: |
|
spec = importlib.util.find_spec(imp) |
|
if spec is None: |
|
log(f"{imp} is required, please check requirements are installed.", |
|
type="ERROR", always=True) |
|
return False |
|
|
|
install_js() |
|
return True |
|
|
|
|
|
def get_async_loop(): |
|
loop = None |
|
try: |
|
loop = asyncio.get_event_loop() |
|
except: |
|
loop = asyncio.new_event_loop() |
|
asyncio.set_event_loop(loop) |
|
return loop |
|
|
|
|
|
def get_http_session(): |
|
loop = get_async_loop() |
|
return aiohttp.ClientSession(loop=loop) |
|
|
|
|
|
async def download(url, stream, update_callback=None, session=None): |
|
close_session = False |
|
if session is None: |
|
close_session = True |
|
session = get_http_session() |
|
try: |
|
async with session.get(url) as response: |
|
size = int(response.headers.get('content-length', 0)) or None |
|
|
|
with tqdm( |
|
unit='B', unit_scale=True, miniters=1, desc=url.split('/')[-1], total=size, |
|
) as progressbar: |
|
perc = 0 |
|
async for chunk in response.content.iter_chunked(2048): |
|
stream.write(chunk) |
|
progressbar.update(len(chunk)) |
|
if update_callback is not None and progressbar.total is not None and progressbar.total != 0: |
|
last = perc |
|
perc = round(progressbar.n / progressbar.total, 2) |
|
if perc != last: |
|
last = perc |
|
await update_callback(perc) |
|
finally: |
|
if close_session and session is not None: |
|
await session.close() |
|
|
|
|
|
async def download_to_file(url, destination, update_callback=None, is_ext_subpath=True, session=None): |
|
if is_ext_subpath: |
|
destination = get_ext_dir(destination) |
|
with open(destination, mode='wb') as f: |
|
download(url, f, update_callback, session) |
|
|
|
|
|
def wait_for_async(async_fn, loop=None): |
|
res = [] |
|
|
|
async def run_async(): |
|
r = await async_fn() |
|
res.append(r) |
|
|
|
if loop is None: |
|
try: |
|
loop = asyncio.get_event_loop() |
|
except: |
|
loop = asyncio.new_event_loop() |
|
asyncio.set_event_loop(loop) |
|
|
|
loop.run_until_complete(run_async()) |
|
|
|
return res[0] |
|
|
|
|
|
def update_node_status(client_id, node, text, progress=None): |
|
if client_id is None: |
|
client_id = PromptServer.instance.client_id |
|
|
|
if client_id is None: |
|
return |
|
|
|
PromptServer.instance.send_sync("pysssss/update_status", { |
|
"node": node, |
|
"progress": progress, |
|
"text": text |
|
}, client_id) |
|
|
|
|
|
async def update_node_status_async(client_id, node, text, progress=None): |
|
if client_id is None: |
|
client_id = PromptServer.instance.client_id |
|
|
|
if client_id is None: |
|
return |
|
|
|
await PromptServer.instance.send("pysssss/update_status", { |
|
"node": node, |
|
"progress": progress, |
|
"text": text |
|
}, client_id) |
|
|
|
|
|
def get_config_value(key, default=None, throw=False): |
|
split = key.split(".") |
|
obj = get_extension_config() |
|
for s in split: |
|
if s in obj: |
|
obj = obj[s] |
|
else: |
|
if throw: |
|
raise KeyError("Configuration key missing: " + key) |
|
else: |
|
return default |
|
return obj |
|
|
|
|
|
def is_inside_dir(root_dir, check_path): |
|
root_dir = os.path.abspath(root_dir) |
|
if not os.path.isabs(check_path): |
|
check_path = os.path.abspath(os.path.join(root_dir, check_path)) |
|
return os.path.commonpath([check_path, root_dir]) == root_dir |
|
|
|
|
|
def get_child_dir(root_dir, child_path, throw_if_outside=True): |
|
child_path = os.path.abspath(os.path.join(root_dir, child_path)) |
|
if is_inside_dir(root_dir, child_path): |
|
return child_path |
|
if throw_if_outside: |
|
raise NotADirectoryError( |
|
"Saving outside the target folder is not allowed.") |
|
return None |
|
|