634 lines
19 KiB
Python
634 lines
19 KiB
Python
import requests
|
|
import os
|
|
import ast
|
|
import argparse
|
|
import sys
|
|
import subprocess
|
|
from importlib import import_module
|
|
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
|
|
|
|
here = Path(__file__).parent
|
|
executable = Path(sys.executable)
|
|
|
|
# - detect mode
|
|
mode = None
|
|
if os.environ.get("COLAB_GPU"):
|
|
mode = "colab"
|
|
elif "python_embeded" in str(executable):
|
|
mode = "embeded"
|
|
elif ".venv" in str(executable):
|
|
mode = "venv"
|
|
|
|
|
|
if mode is None:
|
|
mode = "unknown"
|
|
|
|
# - Constants
|
|
repo_url = "https://github.com/melmass/comfy_mtb.git"
|
|
repo_owner = "melmass"
|
|
repo_name = "comfy_mtb"
|
|
short_platform = {
|
|
"windows": "win_amd64",
|
|
"linux": "linux_x86_64",
|
|
}
|
|
current_platform = platform.system().lower()
|
|
|
|
# region ansi
|
|
# ANSI escape sequences for text styling
|
|
ANSI_FORMATS = {
|
|
"reset": "\033[0m",
|
|
"bold": "\033[1m",
|
|
"dim": "\033[2m",
|
|
"italic": "\033[3m",
|
|
"underline": "\033[4m",
|
|
"blink": "\033[5m",
|
|
"reverse": "\033[7m",
|
|
"strike": "\033[9m",
|
|
}
|
|
|
|
ANSI_COLORS = {
|
|
"black": "\033[30m",
|
|
"red": "\033[31m",
|
|
"green": "\033[32m",
|
|
"yellow": "\033[33m",
|
|
"blue": "\033[34m",
|
|
"magenta": "\033[35m",
|
|
"cyan": "\033[36m",
|
|
"white": "\033[37m",
|
|
"bright_black": "\033[30;1m",
|
|
"bright_red": "\033[31;1m",
|
|
"bright_green": "\033[32;1m",
|
|
"bright_yellow": "\033[33;1m",
|
|
"bright_blue": "\033[34;1m",
|
|
"bright_magenta": "\033[35;1m",
|
|
"bright_cyan": "\033[36;1m",
|
|
"bright_white": "\033[37;1m",
|
|
"bg_black": "\033[40m",
|
|
"bg_red": "\033[41m",
|
|
"bg_green": "\033[42m",
|
|
"bg_yellow": "\033[43m",
|
|
"bg_blue": "\033[44m",
|
|
"bg_magenta": "\033[45m",
|
|
"bg_cyan": "\033[46m",
|
|
"bg_white": "\033[47m",
|
|
"bg_bright_black": "\033[40;1m",
|
|
"bg_bright_red": "\033[41;1m",
|
|
"bg_bright_green": "\033[42;1m",
|
|
"bg_bright_yellow": "\033[43;1m",
|
|
"bg_bright_blue": "\033[44;1m",
|
|
"bg_bright_magenta": "\033[45;1m",
|
|
"bg_bright_cyan": "\033[46;1m",
|
|
"bg_bright_white": "\033[47;1m",
|
|
}
|
|
|
|
|
|
def apply_format(text, *formats):
|
|
"""Apply ANSI escape sequences for the specified formats to the given text."""
|
|
formatted_text = text
|
|
for format in formats:
|
|
formatted_text = f"{ANSI_FORMATS.get(format, '')}{formatted_text}{ANSI_FORMATS.get('reset', '')}"
|
|
return formatted_text
|
|
|
|
|
|
def apply_color(text, color=None, background=None):
|
|
"""Apply ANSI escape sequences for the specified color and background to the given text."""
|
|
formatted_text = text
|
|
if color:
|
|
formatted_text = f"{ANSI_COLORS.get(color, '')}{formatted_text}{ANSI_FORMATS.get('reset', '')}"
|
|
if background:
|
|
formatted_text = f"{ANSI_COLORS.get(background, '')}{formatted_text}{ANSI_FORMATS.get('reset', '')}"
|
|
return formatted_text
|
|
|
|
|
|
def print_formatted(text, *formats, color=None, background=None, **kwargs):
|
|
"""Print the given text with the specified formats, color, and background."""
|
|
formatted_text = apply_format(text, *formats)
|
|
formatted_text = apply_color(formatted_text, color, background)
|
|
file = kwargs.get("file", sys.stdout)
|
|
header = "[mtb install] "
|
|
|
|
# Handle console encoding for Unicode characters (utf-8)
|
|
encoded_header = header.encode(sys.stdout.encoding, errors="replace").decode(
|
|
sys.stdout.encoding
|
|
)
|
|
encoded_text = formatted_text.encode(sys.stdout.encoding, errors="replace").decode(
|
|
sys.stdout.encoding
|
|
)
|
|
|
|
print(
|
|
" " * len(encoded_header)
|
|
if kwargs.get("no_header")
|
|
else apply_color(apply_format(encoded_header, "bold"), color="yellow"),
|
|
encoded_text,
|
|
file=file,
|
|
)
|
|
|
|
|
|
# endregion
|
|
|
|
|
|
# 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 = []
|
|
|
|
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))} "
|
|
else:
|
|
raise ValueError(
|
|
"Invalid 'cmd' argument. It must be a string or a list of arguments."
|
|
)
|
|
|
|
process = subprocess.Popen(
|
|
shell_cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
universal_newlines=True,
|
|
shell=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()
|
|
|
|
interrupted = False
|
|
|
|
def signal_handler(signum, frame):
|
|
nonlocal interrupted
|
|
interrupted = True
|
|
print("Command execution interrupted.")
|
|
|
|
# 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}")
|
|
|
|
|
|
# endregion
|
|
|
|
try:
|
|
import requirements
|
|
except ImportError:
|
|
print_formatted("Installing requirements-parser...", "italic", color="yellow")
|
|
run_command([executable, "-m", "pip", "install", "requirements-parser"])
|
|
import requirements
|
|
|
|
print_formatted("Done.", "italic", color="green")
|
|
|
|
try:
|
|
from tqdm import tqdm
|
|
except ImportError:
|
|
print_formatted("Installing tqdm...", "italic", color="yellow")
|
|
run_command([executable, "-m", "pip", "install", "--upgrade", "tqdm"])
|
|
from tqdm import tqdm
|
|
|
|
pip_map = {
|
|
"onnxruntime-gpu": "onnxruntime",
|
|
"opencv-contrib": "cv2",
|
|
"tb-nightly": "tensorboard",
|
|
"protobuf": "google.protobuf",
|
|
# Add more mappings as needed
|
|
}
|
|
|
|
|
|
def is_pipe():
|
|
if not sys.stdin.isatty():
|
|
return False
|
|
if sys.platform == "win32":
|
|
try:
|
|
import msvcrt
|
|
|
|
return msvcrt.get_osfhandle(0) != -1
|
|
except ImportError:
|
|
return False
|
|
else:
|
|
try:
|
|
mode = os.fstat(0).st_mode
|
|
return (
|
|
stat.S_ISFIFO(mode)
|
|
or stat.S_ISREG(mode)
|
|
or stat.S_ISBLK(mode)
|
|
or stat.S_ISSOCK(mode)
|
|
)
|
|
except OSError:
|
|
return False
|
|
|
|
|
|
@contextmanager
|
|
def suppress_std():
|
|
with open(os.devnull, "w") as devnull:
|
|
old_stdout = sys.stdout
|
|
old_stderr = sys.stderr
|
|
sys.stdout = devnull
|
|
sys.stderr = devnull
|
|
|
|
try:
|
|
yield
|
|
finally:
|
|
sys.stdout = old_stdout
|
|
sys.stderr = old_stderr
|
|
|
|
|
|
# Get the version from __init__.py
|
|
def get_local_version():
|
|
init_file = os.path.join(os.path.dirname(__file__), "__init__.py")
|
|
if os.path.isfile(init_file):
|
|
with open(init_file, "r") as f:
|
|
tree = ast.parse(f.read())
|
|
for node in ast.walk(tree):
|
|
if isinstance(node, ast.Assign):
|
|
for target in node.targets:
|
|
if (
|
|
isinstance(target, ast.Name)
|
|
and target.id == "__version__"
|
|
and isinstance(node.value, ast.Str)
|
|
):
|
|
return node.value.s
|
|
return None
|
|
|
|
|
|
def download_file(url, file_name):
|
|
with requests.get(url, stream=True) as response:
|
|
response.raise_for_status()
|
|
total_size = int(response.headers.get("content-length", 0))
|
|
with open(file_name, "wb") as file, tqdm(
|
|
desc=file_name.stem,
|
|
total=total_size,
|
|
unit="B",
|
|
unit_scale=True,
|
|
unit_divisor=1024,
|
|
) as progress_bar:
|
|
for chunk in response.iter_content(chunk_size=8192):
|
|
file.write(chunk)
|
|
progress_bar.update(len(chunk))
|
|
|
|
|
|
def get_requirements(path: Path):
|
|
with open(path.resolve(), "r") as requirements_file:
|
|
requirements_txt = requirements_file.read()
|
|
|
|
try:
|
|
parsed_requirements = requirements.parse(requirements_txt)
|
|
except AttributeError:
|
|
print_formatted(
|
|
f"Failed to parse {path}. Please make sure the file is correctly formatted.",
|
|
"bold",
|
|
color="red",
|
|
)
|
|
|
|
return
|
|
|
|
return parsed_requirements
|
|
|
|
|
|
def try_import(requirement):
|
|
dependency = requirement.name.strip()
|
|
import_name = pip_map.get(dependency, dependency)
|
|
installed = False
|
|
|
|
pip_name = dependency
|
|
pip_spec = "".join(specs[0]) if (specs := requirement.specs) else ""
|
|
try:
|
|
with suppress_std():
|
|
import_module(import_name)
|
|
print_formatted(
|
|
f"\t✅ Package {pip_name} already installed (import name: '{import_name}').",
|
|
"bold",
|
|
color="green",
|
|
no_header=True,
|
|
)
|
|
installed = True
|
|
except ImportError:
|
|
print_formatted(
|
|
f"\t⛔ Package {pip_name} is missing (import name: '{import_name}').",
|
|
"bold",
|
|
color="red",
|
|
no_header=True,
|
|
)
|
|
|
|
return (installed, pip_name, pip_spec, import_name)
|
|
|
|
|
|
def import_or_install(requirement, dry=False):
|
|
installed, pip_name, pip_spec, import_name = try_import(requirement)
|
|
|
|
pip_install_name = pip_name + pip_spec
|
|
|
|
if not installed:
|
|
print_formatted(f"Installing package {pip_name}...", "italic", color="yellow")
|
|
if dry:
|
|
print_formatted(
|
|
f"Dry-run: Package {pip_install_name} would be installed (import name: '{import_name}').",
|
|
color="yellow",
|
|
)
|
|
else:
|
|
try:
|
|
run_command([executable, "-m", "pip", "install", pip_install_name])
|
|
print_formatted(
|
|
f"Package {pip_install_name} installed successfully using pip package name (import name: '{import_name}')",
|
|
"bold",
|
|
color="green",
|
|
)
|
|
except subprocess.CalledProcessError as e:
|
|
print_formatted(
|
|
f"Failed to install package {pip_install_name} using pip package name (import name: '{import_name}'). Error: {str(e)}",
|
|
"bold",
|
|
color="red",
|
|
)
|
|
|
|
|
|
def get_github_assets(tag=None):
|
|
if tag:
|
|
tag_url = (
|
|
f"https://api.github.com/repos/{repo_owner}/{repo_name}/releases/tags/{tag}"
|
|
)
|
|
else:
|
|
tag_url = (
|
|
f"https://api.github.com/repos/{repo_owner}/{repo_name}/releases/latest"
|
|
)
|
|
response = requests.get(tag_url)
|
|
if response.status_code == 404:
|
|
# print_formatted(
|
|
# f"Tag version '{apply_color(version,'cyan')}' not found for {owner}/{repo} repository."
|
|
# )
|
|
print_formatted("Error retrieving the release assets.", color="red")
|
|
sys.exit()
|
|
|
|
tag_data = response.json()
|
|
tag_name = tag_data["name"]
|
|
|
|
return tag_data, tag_name
|
|
|
|
|
|
# Install dependencies from requirements.txt
|
|
def install_dependencies(dry=False):
|
|
parsed_requirements = get_requirements(here / "reqs.txt")
|
|
if not parsed_requirements:
|
|
return
|
|
print_formatted(
|
|
"Installing dependencies from reqs.txt...", "italic", color="yellow"
|
|
)
|
|
|
|
for requirement in parsed_requirements:
|
|
import_or_install(requirement, dry=dry)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
full = False
|
|
if len(sys.argv) == 1:
|
|
print_formatted(
|
|
"No arguments provided, doing a full install/update...",
|
|
"italic",
|
|
color="yellow",
|
|
)
|
|
|
|
full = True
|
|
|
|
# Parse command-line arguments
|
|
parser = argparse.ArgumentParser(description="Comfy_mtb install script")
|
|
parser.add_argument(
|
|
"--path",
|
|
"-p",
|
|
type=str,
|
|
help="Path to clone the repository to (i.e the absolute path to ComfyUI/custom_nodes)",
|
|
)
|
|
parser.add_argument(
|
|
"--wheels", "-w", action="store_true", help="Install wheel dependencies"
|
|
)
|
|
parser.add_argument(
|
|
"--requirements", "-r", action="store_true", help="Install requirements.txt"
|
|
)
|
|
parser.add_argument(
|
|
"--dry",
|
|
action="store_true",
|
|
help="Print what will happen without doing it (still making requests to the GH Api)",
|
|
)
|
|
|
|
# - keep
|
|
# parser.add_argument(
|
|
# "--version",
|
|
# default=get_local_version(),
|
|
# help="Version to check against the GitHub API",
|
|
# )
|
|
print_formatted("mtb install", "bold", color="yellow")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# wheels_directory = here / "wheels"
|
|
print_formatted(f"Detected environment: {apply_color(mode,'cyan')}")
|
|
|
|
if args.path:
|
|
clone_dir = Path(args.path)
|
|
if not clone_dir.exists():
|
|
print_formatted(
|
|
"The path provided does not exist on disk... It must be pointing to ComfyUI's custom_nodes directory"
|
|
)
|
|
sys.exit()
|
|
|
|
else:
|
|
repo_dir = clone_dir / repo_name
|
|
if not repo_dir.exists():
|
|
print_formatted(f"Cloning to {repo_dir}...", "italic", color="yellow")
|
|
run_command(["git", "clone", "--recursive", repo_url, repo_dir])
|
|
else:
|
|
print_formatted(
|
|
f"Directory {repo_dir} already exists, we will update it..."
|
|
)
|
|
run_command(["git", "pull", "-C", repo_dir])
|
|
# os.chdir(clone_dir)
|
|
here = clone_dir
|
|
full = True
|
|
|
|
# Install dependencies from requirements.txt
|
|
# if args.requirements or mode == "venv":
|
|
|
|
# if (not args.wheels and mode not in ["colab", "embeded"]) and not full:
|
|
# print_formatted(
|
|
# "Skipping wheel installation. Use --wheels to install wheel dependencies. (only needed for Comfy embed)",
|
|
# "italic",
|
|
# color="yellow",
|
|
# )
|
|
|
|
# install_dependencies(dry=args.dry)
|
|
# sys.exit()
|
|
|
|
# if mode in ["colab", "embeded"]:
|
|
# print_formatted(
|
|
# f"Downloading and installing release wheels since we are in a Comfy {apply_color(mode,'cyan')} environment",
|
|
# "italic",
|
|
# color="yellow",
|
|
# )
|
|
# if full:
|
|
# print_formatted(
|
|
# f"Downloading and installing release wheels since no arguments where provided",
|
|
# "italic",
|
|
# color="yellow",
|
|
# )
|
|
|
|
print_formatted("Checking environment...", "italic", color="yellow")
|
|
missing_deps = []
|
|
if parsed_requirements := get_requirements(here / "reqs.txt"):
|
|
for requirement in parsed_requirements:
|
|
installed, pip_name, pip_spec, import_name = try_import(requirement)
|
|
if not installed:
|
|
missing_deps.append(pip_name.split("-")[0])
|
|
|
|
if not missing_deps:
|
|
print_formatted(
|
|
"All requirements are already installed. Enjoy 🚀",
|
|
"italic",
|
|
color="green",
|
|
)
|
|
sys.exit()
|
|
|
|
# # - Get the tag version from the GitHub API
|
|
# tag_data, tag_name = get_github_assets(tag=None)
|
|
|
|
# # - keep
|
|
# version = args.version
|
|
# # Compare the local and tag versions
|
|
# if version and tag_name:
|
|
# if re.match(r"v?(\d+(\.\d+)+)", version) and re.match(
|
|
# r"v?(\d+(\.\d+)+)", tag_name
|
|
# ):
|
|
# version_parts = [int(part) for part in version.lstrip("v").split(".")]
|
|
# tag_version_parts = [int(part) for part in tag_name.lstrip("v").split(".")]
|
|
|
|
# if version_parts > tag_version_parts:
|
|
# print_formatted(
|
|
# f"Local version ({version}) is greater than the release version ({tag_name}).",
|
|
# "bold",
|
|
# "yellow",
|
|
# )
|
|
# sys.exit()
|
|
|
|
# matching_assets = [
|
|
# asset
|
|
# for asset in tag_data["assets"]
|
|
# if asset["name"].endswith(".whl")
|
|
# and (
|
|
# "any" in asset["name"] or short_platform[current_platform] in asset["name"]
|
|
# )
|
|
# ]
|
|
# if not matching_assets:
|
|
# print_formatted(
|
|
# f"Unsupported operating system: {current_platform}", color="yellow"
|
|
# )
|
|
# wheel_order_asset = next(
|
|
# (asset for asset in tag_data["assets"] if asset["name"] == "wheel_order.txt"),
|
|
# None,
|
|
# )
|
|
# if wheel_order_asset is not None:
|
|
# print_formatted(
|
|
# "⚙️ Sorting the release wheels using wheels order", "italic", color="yellow"
|
|
# )
|
|
# response = requests.get(wheel_order_asset["browser_download_url"])
|
|
# if response.status_code == 200:
|
|
# wheel_order = [line.strip() for line in response.text.splitlines()]
|
|
|
|
# def get_order_index(val):
|
|
# try:
|
|
# return wheel_order.index(val)
|
|
# except ValueError:
|
|
# return len(wheel_order)
|
|
|
|
# matching_assets = sorted(
|
|
# matching_assets,
|
|
# key=lambda x: get_order_index(x["name"].split("-")[0]),
|
|
# )
|
|
# else:
|
|
# print("Failed to fetch wheel_order.txt. Status code:", response.status_code)
|
|
|
|
# missing_deps_urls = []
|
|
# for whl_file in matching_assets:
|
|
# # check if installed
|
|
# missing_deps_urls.append(whl_file["browser_download_url"])
|
|
|
|
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")]
|
|
|
|
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",
|
|
)
|