diff --git a/__init__.py b/__init__.py index dabba90..6b49555 100644 --- a/__init__.py +++ b/__init__.py @@ -13,20 +13,31 @@ import os os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true" os.environ["TF_GPU_ALLOCATOR"] = "cuda_malloc_async" -import traceback -from .log import log, blue_text, cyan_text, get_summary, get_label -from .utils import here -from .utils import comfy_dir -import importlib -import os import ast +import contextlib +import importlib import json +import logging +import os +import shutil +import traceback +from importlib import reload + +from aiohttp import web +from server import PromptServer + +import nodes + +from .endpoint import endlog +from .log import blue_text, cyan_text, get_label, get_summary, log +from .utils import comfy_dir, here NODE_CLASS_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS = {} NODE_CLASS_MAPPINGS_DEBUG = {} +WEB_DIRECTORY = "./web" -__version__ = "0.1.4" +__version__ = "0.2.0" def extract_nodes_from_source(filename): @@ -85,7 +96,7 @@ def load_nodes(): nodes_failed.extend(extract_nodes_from_source(filename)) if errors: - log.info( + log.debug( f"Some nodes failed to load:\n\t" + "\n\t".join(errors) + "\n\n" @@ -100,49 +111,10 @@ def load_nodes(): web_extensions_root = comfy_dir / "web" / "extensions" web_mtb = web_extensions_root / "mtb" -if web_mtb.exists(): - log.debug(f"Web extensions folder found at {web_mtb}") - if not os.path.islink(web_mtb.as_posix()): - log.warn( - f"Web extensions folder at {web_mtb} is not a symlink, if updating please delete it before" - ) +if web_mtb.exists() and hasattr(nodes, "EXTENSION_WEB_DIRS"): + shutil.rmtree(web_mtb) -elif web_extensions_root.exists(): - web_tgt = here / "web" - src = web_tgt.as_posix() - dst = web_mtb.as_posix() - try: - if os.name == "nt": - import _winapi - - _winapi.CreateJunction(src, dst) - else: - os.symlink(web_tgt.as_posix(), web_mtb.as_posix()) - - except OSError: - log.warn(f"Failed to create symlink to {web_mtb}, trying to copy it") - try: - import shutil - - shutil.copytree(web_tgt, web_mtb) - log.info(f"Successfully copied {web_tgt} to {web_mtb}") - except Exception as e: - log.warn( - f"Failed to symlink and copy {web_tgt} to {web_mtb}. Please copy the folder manually." - ) - log.warn(e) - - except Exception as e: - log.warn( - f"Failed to create symlink to {web_mtb}. Please copy the folder manually." - ) - log.warn(e) -else: - log.warn( - f"Comfy root probably not found automatically, please copy the folder {web_mtb} manually in the web/extensions folder of ComfyUI" - ) - # - REGISTER NODES nodes, failed = load_nodes() for node_class in nodes: @@ -166,7 +138,7 @@ for node_class in nodes: ) ) -log.info( +log.debug( f"Loaded the following nodes:\n\t" + "\n\t".join( f"{cyan_text(k)}: {blue_text(get_summary(doc)) if doc else '-'}" @@ -174,19 +146,25 @@ log.info( ) ) +log.info(f"loaded {cyan_text(len(nodes))} nodes successfuly") +if failed: + with contextlib.suppress(Exception): + base_url, port = utils.get_server_info() + log.info( + f"Some nodes ({len(failed)}) could not be loaded. This can be ignored, but go to http://{base_url}:{port}/mtb if you want more information." + ) + + # - ENDPOINT -from server import PromptServer -from .log import log -from aiohttp import web -from importlib import reload -import logging -from .endpoint import endlog + if hasattr(PromptServer, "instance"): restore_deps = ["basicsr"] - swap_deps = ["insightface", "onnxruntime"] - + onnx_deps = ["onnxruntime"] + swap_deps = ["insightface"] + onnx_deps node_dependency_mapping = { + "QrCode": ["qrcode"], + "DeepBump": onnx_deps, "FaceSwap": swap_deps, "LoadFaceSwapModel": swap_deps, "LoadFaceAnalysisModel": restore_deps, diff --git a/endpoint.py b/endpoint.py index c0288a7..be6f3da 100644 --- a/endpoint.py +++ b/endpoint.py @@ -1,18 +1,16 @@ -from .utils import ( - here, - import_install, - styles_dir, - backup_file, -) -from aiohttp import web -from .log import mklog import csv +from aiohttp import web + +from .log import mklog +from .utils import backup_file, here, import_install, reqs_map, run_command, styles_dir endlog = mklog("mtb endpoint") # - ACTIONS import platform +import sys +from pathlib import Path import_install("requirements") @@ -21,21 +19,28 @@ def ACTIONS_installDependency(dependency_names=None): if dependency_names is None: return {"error": "No dependency name provided"} endlog.debug(f"Received Install Dependency request for {dependency_names}") - reqs = [] - if platform.system() == "Windows": - reqs = list(requirements.parse((here / "reqs_windows.txt").read_text())) - else: - reqs = list(requirements.parse((here / "reqs.txt").read_text())) - print([x.specs for x in reqs]) - print( - "\n".join([f"{x.line} {''.join(x.specs[0] if x.specs else '')}" for x in reqs]) - ) - for dependency_name in dependency_names: - for req in reqs: - if req.name == dependency_name: - endlog.debug(f"Dependency {dependency_name} installed") - break - return {"success": True} + # reqs = [] + resolved_names = [reqs_map.get(name, name) for name in dependency_names] + try: + run_command([Path(sys.executable), "-m", "pip", "install"] + resolved_names) + return {"success": True} + + except Exception as e: + return {"error": f"Failed to install dependencies: {e}"} + + # if platform.system() == "Windows": + # reqs = list(requirements.parse((here / "reqs_windows.txt").read_text())) + # else: + # reqs = list(requirements.parse((here / "reqs.txt").read_text())) + # print([x.specs for x in reqs]) + # print( + # "\n".join([f"{x.line} {''.join(x.specs[0] if x.specs else '')}" for x in reqs]) + # ) + # for dependency_name in dependency_names: + # for req in reqs: + # if req.name == dependency_name: + # endlog.debug(f"Dependency {dependency_name} installed") + # break def ACTIONS_getStyles(style_name=None): diff --git a/install.py b/install.py index dc34324..60589d6 100644 --- a/install.py +++ b/install.py @@ -1,20 +1,17 @@ -import requests -import os -import ast import argparse -import sys -import subprocess -from importlib import import_module +import ast +import os import platform -from pathlib import Path -import sys -import stat -import threading -import signal -from contextlib import suppress -from queue import Queue, Empty -from contextlib import contextmanager import shlex +import stat +import subprocess +import sys +from contextlib import contextmanager, suppress +from importlib import import_module +from pathlib import Path +from queue import Empty, Queue + +import requests here = Path(__file__).parent executable = Path(sys.executable) @@ -137,12 +134,6 @@ def print_formatted(text, *formats, color=None, background=None, **kwargs): # region utils -def enqueue_output(out, queue): - for char in iter(lambda: out.read(1), b""): - queue.put(char) - out.close() - - def run_command(cmd, ignored_lines_start=None): if ignored_lines_start is None: ignored_lines_start = [] @@ -150,86 +141,49 @@ def run_command(cmd, ignored_lines_start=None): if isinstance(cmd, str): shell_cmd = cmd elif isinstance(cmd, list): - shell_cmd = "" - for arg in cmd: - if isinstance(arg, Path): - arg = arg.as_posix() - shell_cmd += f"{shlex.quote(str(arg))} " + shell_cmd = " ".join( + arg.as_posix() if isinstance(arg, Path) else shlex.quote(str(arg)) + for arg in cmd + ) else: raise ValueError( "Invalid 'cmd' argument. It must be a string or a list of arguments." ) - process = subprocess.Popen( + try: + _run_command(shell_cmd, ignored_lines_start) + except subprocess.CalledProcessError as e: + print(f"Command failed with return code: {e.returncode}", file=sys.stderr) + print(e.stderr.strip(), file=sys.stderr) + + except KeyboardInterrupt: + print("Command execution interrupted.") + + +def _run_command(shell_cmd, ignored_lines_start): + print_formatted(f"Running {shell_cmd}", "bold") + result = subprocess.run( shell_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, - universal_newlines=True, + text=True, shell=True, + check=True, ) - # Create separate threads to read standard output and standard error streams - stdout_queue = Queue() - stderr_queue = Queue() - stdout_thread = threading.Thread( - target=enqueue_output, args=(process.stdout, stdout_queue) - ) - stderr_thread = threading.Thread( - target=enqueue_output, args=(process.stderr, stderr_queue) - ) - stdout_thread.daemon = True - stderr_thread.daemon = True - stdout_thread.start() - stderr_thread.start() + stdout_lines = result.stdout.strip().split("\n") + stderr_lines = result.stderr.strip().split("\n") - interrupted = False + # Print stdout, skipping ignored lines + for line in stdout_lines: + if not any(line.startswith(ign) for ign in ignored_lines_start): + print(line) - def signal_handler(signum, frame): - nonlocal interrupted - interrupted = True - print("Command execution interrupted.") + # Print stderr + for line in stderr_lines: + print(line, file=sys.stderr) - # Register the signal handler for keyboard interrupts (SIGINT) - signal.signal(signal.SIGINT, signal_handler) - - stdout_buffer = "" - stderr_buffer = "" - - # Process output from both streams until the process completes or interrupted - while not interrupted and ( - process.poll() is None or not stdout_queue.empty() or not stderr_queue.empty() - ): - with suppress(Empty): - stdout_char = stdout_queue.get_nowait() - stdout_buffer += stdout_char - if stdout_char == "\n": - if not any( - stdout_buffer.startswith(ign) for ign in ignored_lines_start - ): - print(stdout_buffer.strip()) - stdout_buffer = "" - with suppress(Empty): - stderr_char = stderr_queue.get_nowait() - stderr_buffer += stderr_char - if stderr_char == "\n": - print(stderr_buffer.strip()) - stderr_buffer = "" - - # Print any remaining content in buffers - if stdout_buffer and not any( - stdout_buffer.startswith(ign) for ign in ignored_lines_start - ): - print(stdout_buffer.strip()) - if stderr_buffer: - print(stderr_buffer.strip()) - - return_code = process.returncode - - if return_code == 0 and not interrupted: - print("Command executed successfully!") - else: - if not interrupted: - print(f"Command failed with return code: {return_code}") + print("Command executed successfully!") # endregion @@ -255,6 +209,8 @@ pip_map = { "opencv-contrib": "cv2", "tb-nightly": "tensorboard", "protobuf": "google.protobuf", + "qrcode[pil]": "qrcode", + "requirements-parser": "requirements" # Add more mappings as needed } @@ -441,16 +397,13 @@ def install_dependencies(dry=False): import_or_install(requirement, dry=dry) -if __name__ == "__main__": +def main(): full = False if len(sys.argv) == 1: print_formatted( - "No arguments provided, doing a full install/update...", - "italic", - color="yellow", + "mtb doesn't need an install script anymore.", "italic", color="yellow" ) - - full = True + return # Parse command-line arguments parser = argparse.ArgumentParser(description="Comfy_mtb install script") @@ -612,22 +565,26 @@ if __name__ == "__main__": # # check if installed # missing_deps_urls.append(whl_file["browser_download_url"]) - install_cmd = [executable, "-m", "pip", "install"] + # install_cmd = [executable, "-m", "pip", "install"] # - Install all deps - if not args.dry: - if platform.system() == "Windows": - wheel_cmd = install_cmd + ["-r", (here / "reqs_windows.txt")] - else: - wheel_cmd = install_cmd + ["-r", (here / "reqs.txt")] + # if not args.dry: + # if platform.system() == "Windows": + # wheel_cmd = install_cmd + ["-r", (here / "reqs_windows.txt")] + # else: + # wheel_cmd = install_cmd + ["-r", (here / "reqs.txt")] - run_command(wheel_cmd) - print_formatted( - "✅ Successfully installed all dependencies.", "italic", color="green" - ) - else: - print_formatted( - f"Would have run the following command:\n\t{apply_color(' '.join(install_cmd),'cyan')}", - "italic", - color="yellow", - ) + # run_command(wheel_cmd) + # print_formatted( + # "✅ Successfully installed all dependencies.", "italic", color="green" + # ) + # else: + # print_formatted( + # f"Would have run the following command:\n\t{apply_color(' '.join(install_cmd),'cyan')}", + # "italic", + # color="yellow", + # ) + + +if __name__ == "__main__": + main() diff --git a/reqs.txt b/reqs.txt deleted file mode 100644 index 852259a..0000000 --- a/reqs.txt +++ /dev/null @@ -1,7 +0,0 @@ -onnxruntime-gpu==1.15.1 -qrcode[pil] -rembg==2.0.50 -tensorflow -facexlib==0.3.0 -insightface==0.7.3 -basicsr==1.4.2 diff --git a/reqs_windows.txt b/reqs_windows.txt deleted file mode 100644 index 4c57cac..0000000 --- a/reqs_windows.txt +++ /dev/null @@ -1,18 +0,0 @@ -https://github.com/melMass/comfy_mtb/releases/download/v0.1.3/pycocotools-2.0.6-cp310-cp310-win_amd64.whl -https://github.com/melMass/comfy_mtb/releases/download/v0.1.3/future-0.18.3-py3-none-any.whl -https://github.com/melMass/comfy_mtb/releases/download/v0.1.3/filterpy-1.4.5-py3-none-any.whl -https://github.com/melMass/comfy_mtb/releases/download/v0.1.3/easydict-1.10-py3-none-any.whl -https://github.com/melMass/comfy_mtb/releases/download/v0.1.3/gdown-4.7.1-py3-none-any.whl -https://github.com/melMass/comfy_mtb/releases/download/v0.1.3/basicsr-1.4.2-py3-none-any.whl -https://github.com/melMass/comfy_mtb/releases/download/v0.1.3/mmcv-2.0.0-py2.py3-none-any.whl -https://github.com/melMass/comfy_mtb/releases/download/v0.1.3/insightface-0.7.3-cp310-cp310-win_amd64.whl - -onnxruntime-gpu==1.15.1 -qrcode[pil] -rembg==2.0.50 -# on windows non WSL 2.10 is the last version with GPU support -tensorflow==2.10.1; -tb-nightly==2.12.0a20230126; platform_system == "Windows" -facexlib==0.3.0 -# the old tf version on windows comes with a breaking protobuf version -protobuf==3.19.6 \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..6cdd271 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,6 @@ +qrcode[pil] +onnxruntime-gpu +requirements-parser +# opencv-contrib +rembg +imageio_ffmpeg \ No newline at end of file diff --git a/utils.py b/utils.py index 0adb3f6..85c05c3 100644 --- a/utils.py +++ b/utils.py @@ -1,6 +1,7 @@ import functools import math import os +import shlex import shutil import signal import socket @@ -19,6 +20,8 @@ import requests import torch from PIL import Image +from .install import pip_map + try: from .log import log except ImportError: @@ -147,102 +150,76 @@ def add_path(path, prepend=False): sys.path.append(path) -def enqueue_output(out, queue): - for line in iter(out.readline, b""): - queue.put(line) - out.close() +def run_command(cmd, ignored_lines_start=None): + if ignored_lines_start is None: + ignored_lines_start = [] - -def run_command(cmd): if isinstance(cmd, str): shell_cmd = cmd elif isinstance(cmd, list): - shell_cmd = "" - for arg in cmd: - if isinstance(arg, Path): - arg = arg.as_posix() - shell_cmd += f"{arg} " + shell_cmd = " ".join( + arg.as_posix() if isinstance(arg, Path) else shlex.quote(str(arg)) + for arg in cmd + ) else: raise ValueError( "Invalid 'cmd' argument. It must be a string or a list of arguments." ) - process = subprocess.Popen( + try: + _run_command(shell_cmd, ignored_lines_start) + except subprocess.CalledProcessError as e: + print(f"Command failed with return code: {e.returncode}", file=sys.stderr) + print(e.stderr.strip(), file=sys.stderr) + + except KeyboardInterrupt: + print("Command execution interrupted.") + + +def _run_command(shell_cmd, ignored_lines_start): + log.debug(f"Running {shell_cmd}") + + result = subprocess.run( shell_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, - universal_newlines=True, + text=True, shell=True, + check=True, ) - # Create separate threads to read standard output and standard error streams - stdout_queue = Queue() - stderr_queue = Queue() - stdout_thread = threading.Thread( - target=enqueue_output, args=(process.stdout, stdout_queue) - ) - stderr_thread = threading.Thread( - target=enqueue_output, args=(process.stderr, stderr_queue) - ) - stdout_thread.daemon = True - stderr_thread.daemon = True - stdout_thread.start() - stderr_thread.start() + stdout_lines = result.stdout.strip().split("\n") + stderr_lines = result.stderr.strip().split("\n") - interrupted = False + # Print stdout, skipping ignored lines + for line in stdout_lines: + if not any(line.startswith(ign) for ign in ignored_lines_start): + print(line) - def signal_handler(signum, frame): - nonlocal interrupted - interrupted = True - print("Command execution interrupted.") + # Print stderr + for line in stderr_lines: + print(line, file=sys.stderr) - # Register the signal handler for keyboard interrupts (SIGINT) - signal.signal(signal.SIGINT, signal_handler) - - # Process output from both streams until the process completes or interrupted - while not interrupted and ( - process.poll() is None or not stdout_queue.empty() or not stderr_queue.empty() - ): - with suppress(Empty): - stdout_line = stdout_queue.get_nowait() - if stdout_line.strip() != "": - print(stdout_line.strip()) - with suppress(Empty): - stderr_line = stderr_queue.get_nowait() - if stderr_line.strip() != "": - print(stderr_line.strip()) - return_code = process.returncode - - if return_code == 0 and not interrupted: - print("Command executed successfully!") - else: - if not interrupted: - print(f"Command failed with return code: {return_code}") + print("Command executed successfully!") # todo use the requirements library -reqs_map = { - "onnxruntime": "onnxruntime-gpu==1.15.1", - "basicsr": "basicsr==1.4.2", - "rembg": "rembg==2.0.50", - "qrcode": "qrcode[pil]", - "requirements": "requirements-parser==0.5.0", -} +reqs_map = {value: key for key, value in pip_map.items()} + +import importlib def import_install(package_name): - from pip._internal import main as pip_main + package_spec = reqs_map.get(package_name, package_name) try: - __import__(package_name) - except ImportError: - package_spec = reqs_map.get(package_name) - if package_spec is None: - print(f"Installing {package_name}") - package_spec = package_name + importlib.import_module(package_name) - pip_main(["install", package_spec]) - __import__(package_name) + except Exception: # (ImportError, ModuleNotFoundError): + run_command( + [Path(sys.executable).as_posix(), "-m", "pip", "install", package_spec] + ) + importlib.import_module(package_name) # endregion