diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index 1c184f5cd..8b1c826da 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -8,17 +8,18 @@ }, "runArgs": [ "--gpus=all" - ], + ], // Features to add to the dev container. More info: https://containers.dev/features. // Configure tool-specific properties. "customizations": { "vscode": { "settings": { - "python.defaultInterpreterPath": "/workspace/miniconda3/envs/comfystream/bin/python", - "python.venvPath": "/workspace/miniconda3/envs", + "python.defaultInterpreterPath": "/workspace/.venv/bin/python", + "python.venvPath": "/workspace", "python.terminal.activateEnvInCurrentTerminal": false, "python.terminal.activateEnvironment": false, - "terminal.integrated.shellIntegration.enabled": true + "terminal.integrated.shellIntegration.enabled": true, + "python.analysis.autoImportCompletions": true }, "extensions": [ "ms-python.python", diff --git a/.devcontainer/post-create.sh b/.devcontainer/post-create.sh index 84787613b..956fa17f8 100755 --- a/.devcontainer/post-create.sh +++ b/.devcontainer/post-create.sh @@ -5,7 +5,7 @@ cd /workspace/comfystream # Install Comfystream in editable mode. echo -e "\e[32mInstalling Comfystream in editable mode...\e[0m" -/workspace/miniconda3/envs/comfystream/bin/python3 -m pip install -e . --root-user-action=ignore > /dev/null +python -m pip install -e . --root-user-action=ignore > /dev/null # Install npm packages if needed if [ ! -d "/workspace/comfystream/ui/node_modules" ]; then diff --git a/.dockerignore b/.dockerignore index 9bb4397ac..f7752f3eb 100644 --- a/.dockerignore +++ b/.dockerignore @@ -5,3 +5,4 @@ .venv/ .github/ .dockerignore +ui/node_modules diff --git a/.vscode/launch.json b/.vscode/launch.json index 4d442c585..e424c365e 100755 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -16,7 +16,7 @@ "--disable-cuda-malloc", "--front-end-version", "Comfy-Org/ComfyUI_frontend@v1.24.2" ], - "python": "/workspace/miniconda3/envs/comfystream/bin/python", + "python": "${command:python.interpreterPath}", "justMyCode": true }, { diff --git a/configs/nodes.yaml b/configs/nodes.yaml index 49d422a57..b892c7959 100644 --- a/configs/nodes.yaml +++ b/configs/nodes.yaml @@ -2,8 +2,8 @@ nodes: # Core TensorRT nodes comfyui-tensorrt: name: "ComfyUI TensorRT" - url: "https://github.com/yondonfu/ComfyUI_TensorRT.git" - branch: "quantization_with_controlnet_fixes" + url: "https://github.com/eliteprox/ComfyUI_TensorRT.git" + branch: "patch-1" type: "tensorrt" dependencies: - "tensorrt==10.12.0.36" diff --git a/docker/Dockerfile b/docker/Dockerfile index 0e6baecf7..72c97cc1b 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -2,9 +2,15 @@ ARG BASE_IMAGE=livepeer/comfyui-base:latest FROM ${BASE_IMAGE} -ENV PATH="/workspace/miniconda3/bin:${PATH}" \ +# Ensure Bash is the default shell (inherited from base) +SHELL ["/bin/bash", "-c"] + +ENV PATH="/workspace/.venv/bin:/root/.local/bin:${PATH}" \ NVM_DIR=/root/.nvm \ - NODE_VERSION=18.18.0 + NODE_VERSION=18.18.0 \ + VIRTUAL_ENV="/workspace/.venv" \ + PYTHONPATH="/workspace/.venv/lib/python${PYTHON_VERSION}/site-packages" \ + SHELL="/bin/bash" RUN echo "Using base image: ${BASE_IMAGE}" && \ apt update && \ diff --git a/docker/Dockerfile.base b/docker/Dockerfile.base index c8f6b7ff1..f7dcfe1f5 100644 --- a/docker/Dockerfile.base +++ b/docker/Dockerfile.base @@ -1,21 +1,28 @@ ARG BASE_IMAGE=nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04 \ - CONDA_VERSION=latest \ PYTHON_VERSION=3.12 FROM "${BASE_IMAGE}" -ARG CONDA_VERSION \ - PYTHON_VERSION +ARG PYTHON_VERSION + +# Set Bash as the default shell +SHELL ["/bin/bash", "-c"] ENV DEBIAN_FRONTEND=noninteractive \ - CONDA_VERSION="${CONDA_VERSION}" \ - PATH="/workspace/miniconda3/bin:${PATH}" \ - PYTHON_VERSION="${PYTHON_VERSION}" + PYTHON_VERSION="${PYTHON_VERSION}" \ + SHELL="/bin/bash" \ + UV_PYTHON="${PYTHON_VERSION}" \ + UV_COMPILE_BYTECODE=1 \ + UV_LINK_MODE=copy \ + PATH="/workspace/.venv/bin:/root/.local/bin:/usr/local/bin:${PATH}" \ + VIRTUAL_ENV="/workspace/.venv" \ + PYTHONPATH="/workspace/.venv/lib/python${PYTHON_VERSION}/site-packages" # System dependencies RUN apt update && apt install -yqq --no-install-recommends \ git \ wget \ + curl \ nano \ socat \ libsndfile1 \ @@ -27,6 +34,11 @@ RUN apt update && apt install -yqq --no-install-recommends \ swig \ libprotobuf-dev \ protobuf-compiler \ + ffmpeg \ + python3 \ + python3-pip \ + python3-venv \ + curl \ && rm -rf /var/lib/apt/lists/* #enable opengl support with nvidia gpu @@ -38,59 +50,52 @@ RUN printf '%s\n' \ ' }' \ '}' > /usr/share/glvnd/egl_vendor.d/10_nvidia.json -# Conda setup +# Install uv package manager +RUN curl -LsSf https://astral.sh/uv/install.sh | sh + +# Install Python and setup virtual environment using uv RUN mkdir -p /workspace/comfystream && \ - wget "https://repo.anaconda.com/miniconda/Miniconda3-${CONDA_VERSION}-Linux-x86_64.sh" -O /tmp/miniconda.sh && \ - bash /tmp/miniconda.sh -b -p /workspace/miniconda3 && \ - eval "$(/workspace/miniconda3/bin/conda shell.bash hook)" && \ - /workspace/miniconda3/bin/conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main && \ - /workspace/miniconda3/bin/conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r && \ - conda create -n comfystream python="${PYTHON_VERSION}" ffmpeg=6 -c conda-forge -y && \ - rm /tmp/miniconda.sh && echo 'export LD_LIBRARY_PATH=/workspace/miniconda3/envs/comfystream/lib:$LD_LIBRARY_PATH' >> ~/.bashrc + uv python install "${PYTHON_VERSION}" && \ + cd /workspace && \ + uv venv .venv --python="${PYTHON_VERSION}" -RUN conda run -n comfystream --no-capture-output pip install --upgrade pip && \ -conda run -n comfystream --no-capture-output pip install wheel +# Ensure pip exists inside the venv for tools that shell out to `python -m pip` +RUN /workspace/.venv/bin/python -m ensurepip --upgrade && \ + /workspace/.venv/bin/python -m pip install -U pip setuptools wheel && \ + /workspace/.venv/bin/python -m pip --version + +# Install comfy-cli using uv +RUN uv pip install comfy-cli # Copy only files needed for setup COPY ./src/comfystream/scripts /workspace/comfystream/src/comfystream/scripts COPY ./configs /workspace/comfystream/configs -# Clone ComfyUI -RUN git clone --branch v0.3.56 --depth 1 https://github.com/comfyanonymous/ComfyUI.git /workspace/ComfyUI +RUN mkdir -p /workspace + +# Install ComfyUI using comfy-cli +RUN comfy --skip-prompt --workspace /workspace/ComfyUI install --skip-torch-or-directml --nvidia # Copy ComfyStream files into ComfyUI COPY . /workspace/comfystream -RUN conda run -n comfystream --cwd /workspace/comfystream --no-capture-output pip install -r ./src/comfystream/scripts/constraints.txt +RUN cd /workspace/comfystream && uv pip install -r ./src/comfystream/scripts/overrides.txt # Copy comfystream and example workflows to ComfyUI COPY ./workflows/comfyui/* /workspace/ComfyUI/user/default/workflows/ COPY ./test/example-512x512.png /workspace/ComfyUI/input -# Install ComfyUI requirements -RUN conda run -n comfystream --no-capture-output --cwd /workspace/ComfyUI pip install -r requirements.txt --root-user-action=ignore - # Install ComfyStream requirements RUN ln -s /workspace/comfystream /workspace/ComfyUI/custom_nodes/comfystream -RUN conda run -n comfystream --no-capture-output --cwd /workspace/comfystream pip install -e . --root-user-action=ignore -RUN conda run -n comfystream --no-capture-output --cwd /workspace/comfystream python install.py --workspace /workspace/ComfyUI +RUN cd /workspace/comfystream && uv pip install -e . +RUN cd /workspace/comfystream && python install.py --workspace /workspace/ComfyUI # Accept a build-arg that lets CI force-invalidate setup_nodes.py ARG CACHEBUST=static ENV CACHEBUST=${CACHEBUST} -# Run setup_nodes -RUN conda run -n comfystream --no-capture-output --cwd /workspace/comfystream python src/comfystream/scripts/setup_nodes.py --workspace /workspace/ComfyUI - -RUN conda run -n comfystream --no-capture-output pip install "numpy<2.0.0" - -RUN conda run -n comfystream --no-capture-output pip install --no-cache-dir xformers==0.0.30 --no-deps - -# Configure no environment activation by default -RUN conda config --set auto_activate_base false && \ - conda init bash -# Set comfystream environment as default -RUN echo "conda activate comfystream" >> ~/.bashrc +# Ensure bash is used by default for interactive shells +ENV BASH_ENV=/root/.bashrc WORKDIR /workspace/comfystream diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh index ef55c77e2..47551fc4e 100755 --- a/docker/entrypoint.sh +++ b/docker/entrypoint.sh @@ -1,7 +1,6 @@ #!/bin/bash set -e -eval "$(conda shell.bash hook)" # Add help command to show usage show_help() { @@ -76,8 +75,7 @@ fi if [ "$1" = "--download-models" ]; then cd /workspace/comfystream - conda activate comfystream - python src/comfystream/scripts/setup_models.py --workspace /workspace/ComfyUI + python -m comfystream.scripts.setup_models --workspace /workspace/ComfyUI shift fi @@ -89,7 +87,6 @@ FASTERLIVEPORTRAIT_DIR="/workspace/ComfyUI/models/liveportrait_onnx" if [ "$1" = "--build-engines" ]; then cd /workspace/comfystream - conda activate comfystream # Build Static Engine for Dreamshaper - Square (512x512) python src/comfystream/scripts/build_trt.py --model /workspace/ComfyUI/models/unet/dreamshaper-8-dmd-1kstep.safetensors --out-engine /workspace/ComfyUI/output/tensorrt/static-dreamshaper8_SD15_\$stat-b-1-h-512-w-512_00001_.engine --width 512 --height 512 @@ -152,8 +149,7 @@ fi if [ "$1" = "--opencv-cuda" ]; then cd /workspace/comfystream - conda activate comfystream - + # Check if OpenCV CUDA build already exists if [ ! -f "/workspace/comfystream/opencv-cuda-release.tar.gz" ]; then # Download and extract OpenCV CUDA build @@ -176,18 +172,18 @@ if [ "$1" = "--opencv-cuda" ]; then libswscale-dev # Remove existing cv2 package - SITE_PACKAGES_DIR="/workspace/miniconda3/envs/comfystream/lib/python3.12/site-packages" + SITE_PACKAGES_DIR="$(uv python dir --bin)/lib/python3.12/site-packages" rm -rf "${SITE_PACKAGES_DIR}/cv2"* # Copy new cv2 package cp -r /workspace/comfystream/cv2 "${SITE_PACKAGES_DIR}/" # Handle library dependencies - CONDA_ENV_LIB="/workspace/miniconda3/envs/comfystream/lib" - + UV_ENV_LIB="$(uv python dir --bin)/lib" + # Remove existing libstdc++ and copy system one - rm -f "${CONDA_ENV_LIB}/libstdc++.so"* - cp /usr/lib/x86_64-linux-gnu/libstdc++.so* "${CONDA_ENV_LIB}/" + rm -f "${UV_ENV_LIB}/libstdc++.so"* + cp /usr/lib/x86_64-linux-gnu/libstdc++.so* "${UV_ENV_LIB}/" # Copy OpenCV libraries cp /workspace/comfystream/opencv/build/lib/libopencv_* /usr/lib/x86_64-linux-gnu/ @@ -207,20 +203,20 @@ if [ "$START_COMFYUI" = true ] || [ "$START_API" = true ] || [ "$START_UI" = tru # Start supervisord in background /usr/bin/supervisord -c /etc/supervisor/supervisord.conf & sleep 2 # Give supervisord time to start - + # Start requested services if [ "$START_COMFYUI" = true ]; then supervisorctl -c /etc/supervisor/supervisord.conf start comfyui fi - + if [ "$START_API" = true ]; then supervisorctl -c /etc/supervisor/supervisord.conf start comfystream-api fi - + if [ "$START_UI" = true ]; then supervisorctl -c /etc/supervisor/supervisord.conf start comfystream-ui fi - + # Keep the script running tail -f /var/log/supervisord.log fi diff --git a/docker/supervisord.conf b/docker/supervisord.conf index 7b1324b74..84869c213 100644 --- a/docker/supervisord.conf +++ b/docker/supervisord.conf @@ -8,7 +8,7 @@ loglevel=info pidfile=/var/run/supervisord.pid [program:comfyui] -command=bash -c "source /workspace/miniconda3/bin/activate comfystream && python main.py --listen --disable-cuda-malloc --front-end-version Comfy-Org/ComfyUI_frontend@v1.24.2" +command=bash -c "python main.py --listen --disable-cuda-malloc --front-end-version Comfy-Org/ComfyUI_frontend@v1.24.2" directory=/workspace/ComfyUI autostart=false autorestart=true @@ -19,7 +19,7 @@ stdout_logfile_maxbytes=0 loglevel=debug [program:comfystream-api] -command=bash -c "source /workspace/miniconda3/bin/activate comfystream && python app.py --workspace=/workspace/ComfyUI --media-ports=5678,5679,5680 --host=0.0.0.0 --port=8889 --log-level=DEBUG" +command=bash -c "python app.py --workspace=/workspace/ComfyUI --media-ports=5678,5679,5680 --host=0.0.0.0 --port=8889 --log-level=DEBUG" directory=/workspace/comfystream/server autostart=false autorestart=true diff --git a/install.py b/install.py index 385d160fd..501d5978f 100644 --- a/install.py +++ b/install.py @@ -66,7 +66,7 @@ def download_and_extract_ui_files(version: str): if __name__ == "__main__": parser = argparse.ArgumentParser(description="Install custom node requirements") parser.add_argument( - "--workspace", default=os.environ.get('COMFY_UI_WORKSPACE', None), required=False, help="Set Comfy workspace" + "--workspace", default=os.environ.get('COMFYUI_WORKSPACE', None), required=False, help="Set Comfy workspace" ) args = parser.parse_args() diff --git a/pyproject.toml b/pyproject.toml index 50e59935d..33ed4486b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,23 +6,16 @@ build-backend = "setuptools.build_meta" name = "comfystream" description = "Build Live AI Video with ComfyUI" version = "0.1.5" -license = { file = "LICENSE" } -dependencies = [ - "asyncio", - "pytrickle @ git+https://github.com/livepeer/pytrickle.git@de37bea74679fa5db46b656a83c9b7240fc597b6", - "comfyui @ git+https://github.com/hiddenswitch/ComfyUI.git@58622c7e91cb5cc2bca985d713db55e5681ff316", - "aiortc", - "aiohttp", - "aiohttp_cors", - "toml", - "twilio", - "prometheus_client", - "librosa" -] +license = "MIT" +dynamic = ["dependencies"] [project.optional-dependencies] dev = ["pytest", "pytest-cov"] +[project.scripts] +setup-nodes = "comfystream.scripts.setup_nodes:setup_nodes" +setup-models = "comfystream.scripts.setup_models:setup_models" + [project.urls] Repository = "https://github.com/yondonfu/comfystream" diff --git a/requirements.txt b/requirements.txt index 790900bb1..6d2e6dac0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,6 @@ asyncio pytrickle @ git+https://github.com/livepeer/pytrickle.git@de37bea74679fa5db46b656a83c9b7240fc597b6 -comfyui @ git+https://github.com/hiddenswitch/ComfyUI.git@58622c7e91cb5cc2bca985d713db55e5681ff316 +comfyui @ git+https://github.com/hiddenswitch/ComfyUI.git@3de77f6e24901fa0c423f1e298a8be9f89f312cc aiortc aiohttp aiohttp_cors diff --git a/scripts/requirements.txt b/scripts/requirements.txt index 6ea80ee2e..eb7cee7f7 100644 --- a/scripts/requirements.txt +++ b/scripts/requirements.txt @@ -8,4 +8,4 @@ bcrypt rich # Profiler psutil -pynvml +nvidia-ml-py diff --git a/src/comfystream/pipeline.py b/src/comfystream/pipeline.py index 3e4febaf1..e9e40d3ec 100644 --- a/src/comfystream/pipeline.py +++ b/src/comfystream/pipeline.py @@ -23,8 +23,12 @@ class Pipeline: postprocessing, and queue management. """ - def __init__(self, width: int = 512, height: int = 512, - comfyui_inference_log_level: Optional[int] = None, **kwargs): + def __init__(self, + width: int = 512, + height: int = 512, + comfyui_inference_log_level: Optional[int] = None, + blacklist_nodes: List[str] = ["ComfyUI-Manager"], + **kwargs): """Initialize the pipeline with the given configuration. Args: diff --git a/src/comfystream/scripts/__init__.py b/src/comfystream/scripts/__init__.py index 9273e2423..a309192f6 100644 --- a/src/comfystream/scripts/__init__.py +++ b/src/comfystream/scripts/__init__.py @@ -1,7 +1,3 @@ -from . import utils -from utils import get_config_path, load_model_config -from .setup_nodes import run_setup_nodes -from .setup_models import run_setup_models """Setup scripts for ComfyUI streaming server""" diff --git a/src/comfystream/scripts/build_trt.py b/src/comfystream/scripts/build_trt.py index 3ef1ac7c3..2b9a1c525 100644 --- a/src/comfystream/scripts/build_trt.py +++ b/src/comfystream/scripts/build_trt.py @@ -5,9 +5,8 @@ import time import argparse -# Reccomended running from comfystream conda environment +# Recommended running from the workspace directory (uv venv auto-activated) # in devcontainer from the workspace/ directory, or comfystream/ if you've checked out the repo -# $> conda activate comfystream # $> python src/comfystream/scripts/build_trt.py --model /ComfyUI/models/checkpoints/SD1.5/dreamshaper-8.safetensors --out-engine /ComfyUI/output/tensorrt/static-dreamshaper8_SD15_$stat-b-1-h-512-w-512_00001_.engine # Paths path explicitly to use the downloaded comfyUI installation on root diff --git a/src/comfystream/scripts/constraints.txt b/src/comfystream/scripts/overrides.txt similarity index 100% rename from src/comfystream/scripts/constraints.txt rename to src/comfystream/scripts/overrides.txt diff --git a/src/comfystream/scripts/setup_models.py b/src/comfystream/scripts/setup_models.py index 9360a542b..d8ed769c7 100644 --- a/src/comfystream/scripts/setup_models.py +++ b/src/comfystream/scripts/setup_models.py @@ -4,13 +4,19 @@ from tqdm import tqdm import yaml import argparse -from utils import get_config_path, load_model_config +from comfystream.scripts.utils import ( + get_config_path, + load_model_config, + get_default_workspace, + validate_and_prompt_workspace, + setup_workspace_environment +) def parse_args(): parser = argparse.ArgumentParser(description='Setup ComfyUI models') parser.add_argument('--workspace', - default=os.environ.get('COMFY_UI_WORKSPACE', os.path.expanduser('~/comfyui')), - help='ComfyUI workspace directory (default: ~/comfyui or $COMFY_UI_WORKSPACE)') + default=get_default_workspace(), + help='ComfyUI workspace directory (default: ~/comfyui or $COMFYUI_WORKSPACE)') return parser.parse_args() def download_file(url, destination, description=None): @@ -121,9 +127,17 @@ def setup_directories(workspace_dir): def setup_models(): args = parse_args() - workspace_dir = Path(args.workspace) - + workspace_dir = validate_and_prompt_workspace(args.workspace, "setup-models") + + setup_workspace_environment(workspace_dir) setup_directories(workspace_dir) setup_model_files(workspace_dir) -setup_models() + +def main(): + """Entry point for command line usage.""" + setup_models() + + +if __name__ == "__main__": + main() diff --git a/src/comfystream/scripts/setup_nodes.py b/src/comfystream/scripts/setup_nodes.py index 418e55f86..632435403 100755 --- a/src/comfystream/scripts/setup_nodes.py +++ b/src/comfystream/scripts/setup_nodes.py @@ -4,15 +4,20 @@ from pathlib import Path import yaml import argparse -from utils import get_config_path, load_model_config - +from comfystream.scripts.utils import ( + get_config_path, + load_model_config, + get_default_workspace, + validate_and_prompt_workspace, + setup_workspace_environment +) def parse_args(): parser = argparse.ArgumentParser(description="Setup ComfyUI nodes and models") parser.add_argument( - "--workspace", - default=os.environ.get("COMFY_UI_WORKSPACE", Path("~/comfyui").expanduser()), - help="ComfyUI workspace directory (default: ~/comfyui or $COMFY_UI_WORKSPACE)", + "--workspace", + default=get_default_workspace(), + help="ComfyUI workspace directory (default: ~/comfyui or $COMFYUI_WORKSPACE)", ) parser.add_argument( "--pull-branches", @@ -20,15 +25,13 @@ def parse_args(): default=False, help="Update existing nodes to their specified branches", ) + parser.add_argument( + "--nodes", + default="nodes", + help="Node configuration file pattern to use. Examples: 'nodes' (default), 'vision', 'utility', 'nodes-vision.yaml'", + ) return parser.parse_args() - -def setup_environment(workspace_dir): - os.environ["COMFY_UI_WORKSPACE"] = str(workspace_dir) - os.environ["PYTHONPATH"] = str(workspace_dir) - os.environ["CUSTOM_NODES_PATH"] = str(workspace_dir / "custom_nodes") - - def setup_directories(workspace_dir): """Create required directories in the workspace""" # Create base directories @@ -36,7 +39,6 @@ def setup_directories(workspace_dir): custom_nodes_dir = workspace_dir / "custom_nodes" custom_nodes_dir.mkdir(parents=True, exist_ok=True) - def install_custom_nodes(workspace_dir, config_path=None, pull_branches=False): """Install custom nodes based on configuration""" if config_path is None: @@ -54,11 +56,13 @@ def install_custom_nodes(workspace_dir, config_path=None, pull_branches=False): custom_nodes_path.mkdir(parents=True, exist_ok=True) os.chdir(custom_nodes_path) - # Get the absolute path to constraints.txt - constraints_path = Path(__file__).parent / "constraints.txt" - if not constraints_path.exists(): - print(f"Warning: constraints.txt not found at {constraints_path}") - constraints_path = None + # Get the absolute path to overrides.txt and set UV_OVERRIDES env var + overrides_path = Path(__file__).parent / "overrides.txt" + if overrides_path.exists(): + os.environ["UV_OVERRIDES"] = str(overrides_path) + print(f"Set UV_OVERRIDES to {overrides_path}") + else: + print(f"Warning: overrides.txt not found at {overrides_path}") try: for _, node_info in config["nodes"].items(): @@ -90,40 +94,42 @@ def install_custom_nodes(workspace_dir, config_path=None, pull_branches=False): # Install requirements if present requirements_file = node_path / "requirements.txt" if requirements_file.exists(): - pip_cmd = [ - sys.executable, - "-m", + print(f"Installing requirements from {requirements_file}") + uv_cmd = [ + "uv", "pip", "install", "-r", str(requirements_file), ] - if constraints_path and constraints_path.exists(): - pip_cmd.extend(["-c", str(constraints_path)]) - subprocess.run(pip_cmd, check=True) + subprocess.run(uv_cmd, check=True) # Install additional dependencies if specified if "dependencies" in node_info: for dep in node_info["dependencies"]: - pip_cmd = [sys.executable, "-m", "pip", "install", dep] - if constraints_path and constraints_path.exists(): - pip_cmd.extend(["-c", str(constraints_path)]) - subprocess.run(pip_cmd, check=True) + uv_cmd = ["uv", "pip", "install", dep] + subprocess.run(uv_cmd, check=True) print(f"Installed {node_info['name']}") except Exception as e: print(f"Error installing {node_info['name']} {e}") raise e - def setup_nodes(): args = parse_args() - workspace_dir = Path(args.workspace) - - setup_environment(workspace_dir) + workspace_dir = validate_and_prompt_workspace(args.workspace, "setup-nodes") + + setup_workspace_environment(workspace_dir) setup_directories(workspace_dir) - install_custom_nodes(workspace_dir, pull_branches=args.pull_branches) + + # Get config path from argument + config_path = get_config_path(args.nodes) + install_custom_nodes(workspace_dir, config_path=config_path, pull_branches=args.pull_branches) -if __name__ == "__main__": +def main(): + """Entry point for command line usage.""" setup_nodes() + +if __name__ == "__main__": + main() diff --git a/src/comfystream/scripts/utils.py b/src/comfystream/scripts/utils.py index a7b37f2df..1eb982fe8 100644 --- a/src/comfystream/scripts/utils.py +++ b/src/comfystream/scripts/utils.py @@ -1,21 +1,82 @@ import yaml +import os from pathlib import Path def get_config_path(filename): - """Get the absolute path to a config file""" - config_path = Path("configs") / filename - if not config_path.exists(): - print(f"Warning: Config file {filename} not found at {config_path}") - print(f"Available files in configs/:") - try: - for f in Path("configs").glob("*"): - print(f" - {f.name}") - except FileNotFoundError: - print(" configs/ directory not found") - raise FileNotFoundError(f"Config file {filename} not found at {config_path}") - return config_path + """Get the absolute path to a config file with pattern matching support""" + configs_dir = Path("configs") + + if not configs_dir.exists(): + print(" configs/ directory not found") + raise FileNotFoundError("configs/ directory not found") + + # First try exact match + config_path = configs_dir / filename + if config_path.exists(): + return config_path + + # If no extension provided, try adding .yaml + if not filename.endswith('.yaml'): + config_path = configs_dir / f"{filename}.yaml" + if config_path.exists(): + return config_path + + # Try pattern matching for nodes-* files + if not filename.startswith('nodes-') and not filename == 'nodes.yaml': + pattern_path = configs_dir / f"nodes-{filename}.yaml" + if pattern_path.exists(): + return pattern_path + + # If still not found, show available files + print(f"Warning: Config file matching '{filename}' not found") + raise FileNotFoundError(f"Config file matching '{filename}' not found") def load_model_config(config_path): """Load model configuration from YAML file""" with open(config_path, 'r') as f: - return yaml.safe_load(f) \ No newline at end of file + return yaml.safe_load(f) + +def get_default_workspace(): + """Get the default workspace directory""" + return os.environ.get("COMFYUI_WORKSPACE", Path("~/comfyui").expanduser()) + +def validate_and_prompt_workspace(workspace_path, script_name="script"): + """ + Validate workspace directory exists and prompt user to create if it doesn't. + + Args: + workspace_path: Path to workspace directory (str or Path object) + script_name: Name of the calling script for better error messages + + Returns: + Path: Validated workspace directory path + + Raises: + SystemExit: If user cancels workspace creation + """ + workspace_dir = Path(workspace_path) + + # Check if workspace exists, and prompt user if it doesn't + if not workspace_dir.exists(): + print(f"Workspace directory '{workspace_dir}' does not exist.") + + # Check if this is the default workspace (user didn't specify one) + default_workspace = get_default_workspace() + if str(workspace_dir) == str(default_workspace): + print("No workspace was specified and the default workspace doesn't exist.") + + try: + response = input(f"Would you like to create '{workspace_dir}' and continue? (y/N): ").strip().lower() + if response not in ['y', 'yes']: + print(f"{script_name} cancelled.") + raise SystemExit(0) + except (KeyboardInterrupt, EOFError): + print(f"\n{script_name} cancelled.") + raise SystemExit(0) + + return workspace_dir + +def setup_workspace_environment(workspace_dir): + """Setup environment variables for workspace""" + os.environ["COMFYUI_WORKSPACE"] = str(workspace_dir) + os.environ["CUSTOM_NODES_PATH"] = str(workspace_dir / "custom_nodes")