comfy-env 0.1.15__py3-none-any.whl → 0.1.16__py3-none-any.whl

comfy_env/install.py

@@ -5,17 +5,13 @@ import os
 from pathlib import Path
 from typing import Callable, List, Optional, Set, Union
 
-from .config.parser import ComfyEnvConfig, NodeReq, load_config, discover_config
+from .config import ComfyEnvConfig, NodeDependency, load_config, discover_config, CONFIG_FILE_NAME
 
-
-# Environment variable to disable comfy-env isolation
 USE_COMFY_ENV_VAR = "USE_COMFY_ENV"
 
 
 def _is_comfy_env_enabled() -> bool:
-    """Check if isolation is enabled."""
-    val = os.environ.get(USE_COMFY_ENV_VAR, "1").lower()
-    return val not in ("0", "false", "no", "off")
+    return os.environ.get(USE_COMFY_ENV_VAR, "1").lower() not in ("0", "false", "no", "off")
 
 
 def install(
@@ -25,15 +21,11 @@ def install(
     dry_run: bool = False,
 ) -> bool:
     """Install dependencies from comfy-env.toml."""
-    # Auto-discover caller's directory if not provided
     if node_dir is None:
-        frame = inspect.stack()[1]
-        caller_file = frame.filename
-        node_dir = Path(caller_file).parent.resolve()
+        node_dir = Path(inspect.stack()[1].filename).parent.resolve()
 
     log = log_callback or print
 
-    # Load config
     if config is not None:
         config_path = Path(config)
         if not config_path.is_absolute():
@@ -43,32 +35,16 @@ def install(
         cfg = discover_config(node_dir)
 
     if cfg is None:
-        raise FileNotFoundError(
-            f"No comfy-env.toml found in {node_dir}. "
-            "Create comfy-env.toml to define dependencies."
-        )
-
-    # Install apt packages first (Linux only)
-    if cfg.apt_packages:
-        _install_apt_packages(cfg.apt_packages, log, dry_run)
+        raise FileNotFoundError(f"No comfy-env.toml found in {node_dir}")
 
-    # Set persistent env vars (for OpenMP settings, etc.)
-    if cfg.env_vars:
-        _set_persistent_env_vars(cfg.env_vars, log, dry_run)
+    if cfg.apt_packages: _install_apt_packages(cfg.apt_packages, log, dry_run)
+    if cfg.env_vars: _set_persistent_env_vars(cfg.env_vars, log, dry_run)
+    if cfg.node_reqs: _install_node_dependencies(cfg.node_reqs, node_dir, log, dry_run)
 
-    # Install node dependencies
-    if cfg.node_reqs:
-        _install_node_dependencies(cfg.node_reqs, node_dir, log, dry_run)
-
-    # Check if isolation is enabled
     if _is_comfy_env_enabled():
-        # Install everything via pixi (isolated environment)
         _install_via_pixi(cfg, node_dir, log, dry_run)
-
-        # Auto-discover and install isolated subdirectory environments
         _install_isolated_subdirs(node_dir, log, dry_run)
     else:
-        # Install directly to host Python (no isolation)
         log("\n[comfy-env] Isolation disabled (USE_COMFY_ENV=0)")
         _install_to_host_python(cfg, node_dir, log, dry_run)
 
@@ -76,369 +52,191 @@ def install(
     return True
 
 
-def _install_apt_packages(
-    packages: List[str],
-    log: Callable[[str], None],
-    dry_run: bool,
-) -> None:
-    """Install apt packages (Linux only)."""
-    import os
+def _install_apt_packages(packages: List[str], log: Callable[[str], None], dry_run: bool) -> None:
+    from .packages.apt import apt_install
     import platform
-    import shutil
-    import subprocess
-
     if platform.system() != "Linux":
-        log(f"[apt] Skipping apt packages (not Linux)")
-        return
-
-    log(f"\n[apt] Installing {len(packages)} system package(s):")
-    for pkg in packages:
-        log(f"  - {pkg}")
-
-    if dry_run:
-        log("  (dry run - no changes made)")
         return
+    log(f"\n[apt] Installing: {', '.join(packages)}")
+    if not dry_run:
+        apt_install(packages, log)
 
-    # Determine if we need sudo
-    is_root = os.geteuid() == 0
-    has_sudo = shutil.which("sudo") is not None
-    use_sudo = not is_root and has_sudo
-    prefix = ["sudo"] if use_sudo else []
-
-    if not is_root and not has_sudo:
-        log(f"[apt] Warning: No root access. Install manually:")
-        log(f"  sudo apt-get update && sudo apt-get install -y {' '.join(packages)}")
-        return
 
-    # Run apt-get update (suppress output, just show errors)
-    log("[apt] Updating package lists...")
-    result = subprocess.run(
-        prefix + ["apt-get", "update"],
-        capture_output=True,
-        text=True,
-    )
-    if result.returncode != 0:
-        log(f"[apt] Warning: apt-get update failed: {result.stderr.strip()}")
-
-    # Install each package individually (some may not exist on all distros)
-    log("[apt] Installing packages...")
-    installed = []
-    skipped = []
-    for pkg in packages:
-        result = subprocess.run(
-            prefix + ["apt-get", "install", "-y", pkg],
-            capture_output=True,
-            text=True,
-        )
-        if result.returncode == 0:
-            installed.append(pkg)
-            log(f"  [apt] Installed {pkg}")
-        else:
-            skipped.append(pkg)
-            log(f"  [apt] Skipped {pkg} (not available)")
-
-    if installed:
-        log(f"[apt] Installed {len(installed)} package(s)")
-    if skipped:
-        log(f"[apt] Skipped {len(skipped)} unavailable package(s)")
-
-
-def _set_persistent_env_vars(
-    env_vars: dict,
-    log: Callable[[str], None],
-    dry_run: bool,
-) -> None:
-    """Set env vars permanently (survives restarts)."""
-    import os
-    import platform
-    import subprocess
-    from pathlib import Path
+def _set_persistent_env_vars(env_vars: dict, log: Callable[[str], None], dry_run: bool) -> None:
+    import platform, subprocess
+    if not env_vars: return
 
-    if not env_vars:
-        return
+    log(f"\n[env] Setting {len(env_vars)} env var(s)")
+    for k, v in env_vars.items(): log(f"  {k}={v}")
+    if dry_run: return
 
     system = platform.system()
-    log(f"\n[env] Setting {len(env_vars)} persistent environment variable(s)...")
-
-    for key, value in env_vars.items():
-        log(f"  - {key}={value}")
-
-    if dry_run:
-        log("  (dry run - no changes made)")
-        return
-
     if system == "Windows":
-        # Windows: use setx (writes to registry)
-        for key, value in env_vars.items():
-            result = subprocess.run(
-                ["setx", key, value],
-                capture_output=True, text=True
-            )
-            if result.returncode == 0:
-                log(f"  [env] Set {key} (Windows registry)")
-            else:
-                log(f"  [env] Warning: Failed to set {key}: {result.stderr.strip()}")
-        log("[env] Restart terminal/ComfyUI for changes to take effect")
-
-    elif system == "Darwin":  # macOS
-        # macOS: launchctl for GUI apps + zshrc for terminal
-        for key, value in env_vars.items():
-            subprocess.run(["launchctl", "setenv", key, value], capture_output=True)
-            log(f"  [env] Set {key} (launchctl)")
-
-        # Also add to zshrc for terminal (zsh is default on macOS)
+        for k, v in env_vars.items():
+            subprocess.run(["setx", k, v], capture_output=True)
+    elif system == "Darwin":
+        for k, v in env_vars.items():
+            subprocess.run(["launchctl", "setenv", k, v], capture_output=True)
         _add_to_shell_profile(env_vars, log)
-
-    else:  # Linux
+    else:
         _add_to_shell_profile(env_vars, log)
 
 
-def _add_to_shell_profile(
-    env_vars: dict,
-    log: Callable[[str], None],
-) -> None:
-    """Add env vars to shell profile (Linux/macOS)."""
-    import os
-    from pathlib import Path
-
-    # Determine shell profile
+def _add_to_shell_profile(env_vars: dict, log: Callable[[str], None]) -> None:
     shell = os.environ.get("SHELL", "/bin/bash")
-    if "zsh" in shell:
-        rc_file = Path.home() / ".zshrc"
-    else:
-        rc_file = Path.home() / ".bashrc"
-
+    rc_file = Path.home() / (".zshrc" if "zsh" in shell else ".bashrc")
     profile_file = Path.home() / ".comfy-env-profile"
 
-    # Write env vars to our dedicated file
     with open(profile_file, "w") as f:
-        f.write("# Generated by comfy-env - do not edit manually\n")
-        for key, value in env_vars.items():
-            f.write(f'export {key}="{value}"\n')
-    log(f"  [env] Wrote {profile_file}")
+        f.write("# Generated by comfy-env\n")
+        for k, v in env_vars.items():
+            f.write(f'export {k}="{v}"\n')
 
-    # Add source line to shell rc (only once)
     source_line = f'source "{profile_file}"'
     existing = rc_file.read_text() if rc_file.exists() else ""
-
-    if source_line not in existing and str(profile_file) not in existing:
+    if source_line not in existing:
         with open(rc_file, "a") as f:
-            f.write(f'\n# comfy-env environment variables\n')
-            f.write(f'{source_line}\n')
-        log(f"  [env] Added source line to {rc_file}")
-    else:
-        log(f"  [env] Already configured in {rc_file}")
-
-    log("[env] Restart terminal/ComfyUI for changes to take effect")
-
+            f.write(f'\n# comfy-env\n{source_line}\n')
+    log(f"  [env] Wrote {profile_file}")
 
-def _install_node_dependencies(
-    node_reqs: List[NodeReq],
-    node_dir: Path,
-    log: Callable[[str], None],
-    dry_run: bool,
-) -> None:
-    """Install node dependencies (other ComfyUI custom nodes)."""
-    from .nodes import install_node_deps
 
+def _install_node_dependencies(node_reqs: List[NodeDependency], node_dir: Path, log: Callable[[str], None], dry_run: bool) -> None:
+    from .packages.node_dependencies import install_node_dependencies
     custom_nodes_dir = node_dir.parent
     log(f"\nInstalling {len(node_reqs)} node dependencies...")
-
     if dry_run:
         for req in node_reqs:
-            node_path = custom_nodes_dir / req.name
-            status = "exists" if node_path.exists() else "would clone"
-            log(f"  {req.name}: {status}")
+            log(f"  {req.name}: {'exists' if (custom_nodes_dir / req.name).exists() else 'would clone'}")
         return
-
-    visited: Set[str] = {node_dir.name}
-    install_node_deps(node_reqs, custom_nodes_dir, log, visited)
+    install_node_dependencies(node_reqs, custom_nodes_dir, log, {node_dir.name})
 
 
-def _install_via_pixi(
-    cfg: ComfyEnvConfig,
-    node_dir: Path,
-    log: Callable[[str], None],
-    dry_run: bool,
-) -> None:
-    """Install all packages via pixi."""
-    from .pixi import pixi_install
+def _install_via_pixi(cfg: ComfyEnvConfig, node_dir: Path, log: Callable[[str], None], dry_run: bool) -> None:
+    from .packages.pixi import ensure_pixi, get_pixi_python, pixi_clean
+    from .packages.toml_generator import write_pixi_toml
+    from .packages.cuda_wheels import get_wheel_url, CUDA_TORCH_MAP
+    from .detection import get_recommended_cuda_version
+    from .environment.cache import get_central_env_path, write_marker, write_env_metadata, MARKER_FILE, get_cache_dir
+    import shutil, subprocess, sys
 
-    # Count what we're installing
-    cuda_count = len(cfg.cuda_packages)
-
-    # Count from passthrough (pixi-native format)
     deps = cfg.pixi_passthrough.get("dependencies", {})
     pypi_deps = cfg.pixi_passthrough.get("pypi-dependencies", {})
-
-    if cuda_count == 0 and not deps and not pypi_deps:
+    if not cfg.cuda_packages and not deps and not pypi_deps:
         log("No packages to install")
         return
 
     log(f"\nInstalling via pixi:")
-    if cuda_count:
-        log(f"  CUDA packages: {', '.join(cfg.cuda_packages)}")
-    if deps:
-        log(f"  Conda packages: {len(deps)}")
-    if pypi_deps:
-        log(f"  PyPI packages: {len(pypi_deps)}")
-
-    if dry_run:
-        log("\n(dry run - no changes made)")
-        return
+    if cfg.cuda_packages: log(f"  CUDA: {', '.join(cfg.cuda_packages)}")
+    if deps: log(f"  Conda: {len(deps)}")
+    if pypi_deps: log(f"  PyPI: {len(pypi_deps)}")
+    if dry_run: return
+
+    pixi_clean(node_dir, log)
+    (node_dir / ".pixi").mkdir(parents=True, exist_ok=True)
+    (node_dir / ".pixi" / "config.toml").write_text("detached-environments = false\n")
+
+    pixi_path = ensure_pixi(log=log)
+    cuda_version = torch_version = None
+    if cfg.has_cuda and sys.platform != "darwin":
+        cuda_version = get_recommended_cuda_version()
+        if cuda_version:
+            torch_version = CUDA_TORCH_MAP.get(".".join(cuda_version.split(".")[:2]), "2.8")
 
-    pixi_install(cfg, node_dir, log)
+    write_pixi_toml(cfg, node_dir, log)
+    log("Running pixi install...")
+    result = subprocess.run([str(pixi_path), "install"], cwd=node_dir, capture_output=True, text=True)
+    if result.returncode != 0:
+        raise RuntimeError(f"pixi install failed: {result.stderr}")
 
+    if cfg.cuda_packages and cuda_version:
+        log(f"Installing CUDA packages...")
+        python_path = get_pixi_python(node_dir)
+        if not python_path:
+            raise RuntimeError("No Python in pixi env")
 
-def _install_to_host_python(
-    cfg: ComfyEnvConfig,
-    node_dir: Path,
-    log: Callable[[str], None],
-    dry_run: bool,
-) -> None:
-    """Install packages directly to host Python (no isolation)."""
-    import shutil
-    import subprocess
-    import sys
+        result = subprocess.run([str(python_path), "-c", "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')"],
+                               capture_output=True, text=True)
+        py_version = result.stdout.strip() if result.returncode == 0 else f"{sys.version_info.major}.{sys.version_info.minor}"
 
-    from .pixi import CUDA_WHEELS_INDEX, find_wheel_url
-    from .pixi.cuda_detection import get_recommended_cuda_version
+        for package in cfg.cuda_packages:
+            wheel_url = get_wheel_url(package, torch_version, cuda_version, py_version)
+            if not wheel_url:
+                raise RuntimeError(f"No wheel for {package}")
+            log(f"  {package}")
+            result = subprocess.run([str(python_path), "-m", "pip", "install", "--no-deps", "--no-cache-dir", wheel_url],
+                                   capture_output=True, text=True)
+            if result.returncode != 0:
+                raise RuntimeError(f"Failed: {result.stderr}")
+
+    old_env = node_dir / ".pixi" / "envs" / "default"
+    config_path = node_dir / "comfy-env.toml"
+    main_node_dir = node_dir
+    for parent in node_dir.parents:
+        if parent.parent.name == "custom_nodes":
+            main_node_dir = parent
+            break
+
+    central_env = get_central_env_path(main_node_dir, config_path)
+    if old_env.exists():
+        get_cache_dir()
+        if central_env.exists(): shutil.rmtree(central_env)
+        shutil.move(str(old_env), str(central_env))
+        write_marker(config_path, central_env)
+        write_env_metadata(central_env, config_path.parent / MARKER_FILE)
+        shutil.rmtree(node_dir / ".pixi", ignore_errors=True)
+        log(f"Env: {central_env}")
+
+
+def _install_to_host_python(cfg: ComfyEnvConfig, node_dir: Path, log: Callable[[str], None], dry_run: bool) -> None:
+    import shutil, subprocess, sys
+    from .packages.cuda_wheels import get_wheel_url, CUDA_TORCH_MAP
+    from .detection import get_recommended_cuda_version
 
-    # Collect packages to install
     pypi_deps = cfg.pixi_passthrough.get("pypi-dependencies", {})
-    conda_deps = cfg.pixi_passthrough.get("dependencies", {})
-
-    # Warn about conda dependencies (can't install without pixi)
-    # Filter out 'python' and 'pip' which are meta-dependencies
-    real_conda_deps = {k: v for k, v in conda_deps.items() if k not in ("python", "pip")}
-    if real_conda_deps:
-        log(f"\n[warning] Cannot install conda packages without isolation:")
-        for pkg in real_conda_deps:
-            log(f"  - {pkg}")
-        log("  Set USE_COMFY_ENV=1 to enable isolated environments")
-
-    # Nothing to install?
     if not pypi_deps and not cfg.cuda_packages:
         log("No packages to install")
         return
 
-    # Build pip install command
     pip_packages = []
-
-    # Add pypi dependencies
     for pkg, spec in pypi_deps.items():
         if isinstance(spec, str):
-            if spec == "*":
-                pip_packages.append(pkg)
-            else:
-                pip_packages.append(f"{pkg}{spec}")
+            pip_packages.append(pkg if spec == "*" else f"{pkg}{spec}")
         elif isinstance(spec, dict):
-            version = spec.get("version", "*")
             extras = spec.get("extras", [])
-            if extras:
-                pkg_with_extras = f"{pkg}[{','.join(extras)}]"
-            else:
-                pkg_with_extras = pkg
-            if version == "*":
-                pip_packages.append(pkg_with_extras)
-            else:
-                pip_packages.append(f"{pkg_with_extras}{version}")
-
-    log(f"\nInstalling to host Python ({sys.executable}):")
-    if pip_packages:
-        log(f"  PyPI packages: {len(pip_packages)}")
-    if cfg.cuda_packages:
-        log(f"  CUDA packages: {', '.join(cfg.cuda_packages)}")
-
-    if dry_run:
-        if pip_packages:
-            log(f"  Would install: {', '.join(pip_packages)}")
-        log("\n(dry run - no changes made)")
-        return
+            version = spec.get("version", "*")
+            name = f"{pkg}[{','.join(extras)}]" if extras else pkg
+            pip_packages.append(name if version == "*" else f"{name}{version}")
 
-    # Use uv if available, otherwise pip
-    use_uv = shutil.which("uv") is not None
+    log(f"\nInstalling to {sys.executable}")
+    if dry_run: return
 
-    # Install regular PyPI packages
+    use_uv = shutil.which("uv")
     if pip_packages:
-        if use_uv:
-            cmd = ["uv", "pip", "install", "--python", sys.executable] + pip_packages
-        else:
-            cmd = [sys.executable, "-m", "pip", "install"] + pip_packages
-
-        log(f"  Running: {' '.join(cmd[:4])}...")
-        result = subprocess.run(cmd, capture_output=True, text=True)
-        if result.returncode != 0:
-            log(f"  [error] pip install failed: {result.stderr.strip()[:200]}")
-        else:
-            log(f"  Installed {len(pip_packages)} package(s)")
-
-    # Install CUDA packages from cuda-wheels
+        cmd = ["uv", "pip", "install", "--python", sys.executable] + pip_packages if use_uv else [sys.executable, "-m", "pip", "install"] + pip_packages
+        subprocess.run(cmd, capture_output=True)
+
     if cfg.cuda_packages:
         cuda_version = get_recommended_cuda_version()
-        if not cuda_version:
-            log("  [warning] No CUDA detected, skipping CUDA packages")
-            return
-
-        # Get torch version for wheel matching
-        cuda_mm = ".".join(cuda_version.split(".")[:2])
-        from .pixi.core import CUDA_TORCH_MAP
-        torch_version = CUDA_TORCH_MAP.get(cuda_mm, "2.8")
-
+        if not cuda_version: return
+        torch_version = CUDA_TORCH_MAP.get(".".join(cuda_version.split(".")[:2]), "2.8")
         py_version = f"{sys.version_info.major}.{sys.version_info.minor}"
-        log(f"  CUDA {cuda_version}, PyTorch {torch_version}, Python {py_version}")
 
         for package in cfg.cuda_packages:
-            wheel_url = find_wheel_url(package, torch_version, cuda_version, py_version)
-            if not wheel_url:
-                log(f"  [error] No wheel found for {package}")
-                continue
+            wheel_url = get_wheel_url(package, torch_version, cuda_version, py_version)
+            if wheel_url:
+                cmd = ["uv", "pip", "install", "--python", sys.executable, "--no-deps", wheel_url] if use_uv else [sys.executable, "-m", "pip", "install", "--no-deps", wheel_url]
+                subprocess.run(cmd, capture_output=True)
 
-            log(f"  Installing {package}...")
-            if use_uv:
-                cmd = ["uv", "pip", "install", "--python", sys.executable, "--no-deps", wheel_url]
-            else:
-                cmd = [sys.executable, "-m", "pip", "install", "--no-deps", wheel_url]
-
-            result = subprocess.run(cmd, capture_output=True, text=True)
-            if result.returncode != 0:
-                log(f"  [error] Failed to install {package}: {result.stderr.strip()[:200]}")
-            else:
-                log(f"  Installed {package}")
 
-
-def _install_isolated_subdirs(
-    node_dir: Path,
-    log: Callable[[str], None],
-    dry_run: bool,
-) -> None:
-    """Find and install comfy-env.toml in subdirectories."""
-    from .pixi import pixi_install
-    from .config.parser import CONFIG_FILE_NAME
-
-    # Find all comfy-env.toml files in subdirectories (not root)
+def _install_isolated_subdirs(node_dir: Path, log: Callable[[str], None], dry_run: bool) -> None:
     for config_file in node_dir.rglob(CONFIG_FILE_NAME):
-        if config_file.parent == node_dir:
-            continue  # Skip root (already installed)
-
-        sub_dir = config_file.parent
-        relative = sub_dir.relative_to(node_dir)
-
-        log(f"\n[isolated] Installing: {relative}")
-        sub_cfg = load_config(config_file)
+        if config_file.parent == node_dir: continue
+        log(f"\n[isolated] {config_file.parent.relative_to(node_dir)}")
+        if not dry_run:
+            _install_via_pixi(load_config(config_file), config_file.parent, log, dry_run)
 
-        if dry_run:
-            log(f"  (dry run)")
-            continue
 
-        pixi_install(sub_cfg, sub_dir, log)
-
-
-def verify_installation(
-    packages: List[str],
-    log: Callable[[str], None] = print,
-) -> bool:
-    """Verify that packages are importable."""
+def verify_installation(packages: List[str], log: Callable[[str], None] = print) -> bool:
     all_ok = True
     for package in packages:
         import_name = package.replace("-", "_").split("[")[0]

comfy_env/isolation/__init__.py

@@ -1,10 +1,40 @@
 """
-Isolation module for wrapping ComfyUI nodes to run in isolated environments.
+Isolation layer - Process isolation for node execution.
+
+Wraps node classes to execute in isolated subprocess environments.
 """
 
-from .wrap import wrap_isolated_nodes, wrap_nodes
+from .wrap import (
+    wrap_isolated_nodes,
+    wrap_nodes,
+)
+from .workers import (
+    Worker,
+    WorkerError,
+    MPWorker,
+    SubprocessWorker,
+)
+from .tensor_utils import (
+    TensorKeeper,
+    keep_tensor,
+    keep_tensors_recursive,
+    prepare_tensor_for_ipc,
+    prepare_for_ipc_recursive,
+)
 
 __all__ = [
+    # Node wrapping
     "wrap_isolated_nodes",
     "wrap_nodes",
+    # Workers
+    "Worker",
+    "WorkerError",
+    "MPWorker",
+    "SubprocessWorker",
+    # Tensor utilities
+    "TensorKeeper",
+    "keep_tensor",
+    "keep_tensors_recursive",
+    "prepare_tensor_for_ipc",
+    "prepare_for_ipc_recursive",
 ]