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

comfy_env/detection/platform.py

@@ -0,0 +1,70 @@
+"""Platform detection - OS, architecture, platform tags."""
+
+from __future__ import annotations
+
+import platform as platform_module
+import sys
+from dataclasses import dataclass
+
+
+@dataclass
+class PlatformInfo:
+    os_name: str       # linux, windows, darwin
+    arch: str          # x86_64, aarch64, arm64
+    platform_tag: str  # linux_x86_64, win_amd64, macosx_11_0_arm64
+
+
+def detect_platform() -> PlatformInfo:
+    return PlatformInfo(
+        os_name=_get_os_name(),
+        arch=platform_module.machine().lower(),
+        platform_tag=get_platform_tag(),
+    )
+
+
+def _get_os_name() -> str:
+    if sys.platform.startswith('linux'): return 'linux'
+    if sys.platform == 'win32': return 'windows'
+    if sys.platform == 'darwin': return 'darwin'
+    return sys.platform
+
+
+_PLATFORM_TAGS = {
+    ('linux', 'x86_64'): 'linux_x86_64', ('linux', 'amd64'): 'linux_x86_64',
+    ('linux', 'aarch64'): 'linux_aarch64',
+    ('win32', 'amd64'): 'win_amd64', ('win32', 'x86_64'): 'win_amd64',
+    ('darwin', 'arm64'): 'macosx_11_0_arm64',
+    ('darwin', 'x86_64'): 'macosx_10_9_x86_64',
+}
+
+
+def get_platform_tag() -> str:
+    key = (sys.platform if sys.platform != 'linux' else 'linux', platform_module.machine().lower())
+    return _PLATFORM_TAGS.get(key, f'{sys.platform}_{platform_module.machine().lower()}')
+
+
+_PIXI_PLATFORMS = {
+    ('linux', 'x86_64'): 'linux-64', ('linux', 'amd64'): 'linux-64',
+    ('linux', 'aarch64'): 'linux-aarch64',
+    ('windows', 'amd64'): 'win-64', ('windows', 'x86_64'): 'win-64',
+    ('darwin', 'arm64'): 'osx-arm64',
+    ('darwin', 'x86_64'): 'osx-64',
+}
+
+
+def get_pixi_platform() -> str:
+    key = (_get_os_name(), platform_module.machine().lower())
+    return _PIXI_PLATFORMS.get(key, f'{key[0]}-{key[1]}')
+
+
+def get_library_extension() -> str:
+    return {'.dll': 'windows', '.dylib': 'darwin'}.get(_get_os_name(), '.so')
+
+
+def get_executable_suffix() -> str:
+    return '.exe' if _get_os_name() == 'windows' else ''
+
+
+def is_linux() -> bool: return _get_os_name() == 'linux'
+def is_windows() -> bool: return _get_os_name() == 'windows'
+def is_macos() -> bool: return _get_os_name() == 'darwin'

comfy_env/detection/runtime.py

@@ -0,0 +1,103 @@
+"""Runtime environment detection - combines all detection into a single snapshot."""
+
+from __future__ import annotations
+
+import sys
+from dataclasses import dataclass
+from typing import Dict, Optional
+
+from .cuda import detect_cuda_version
+from .gpu import detect_cuda_environment
+from .platform import get_platform_tag, _get_os_name
+
+
+@dataclass
+class RuntimeEnv:
+    """Detected runtime environment for wheel resolution."""
+    os_name: str
+    platform_tag: str
+    python_version: str
+    python_short: str
+    cuda_version: Optional[str]
+    cuda_short: Optional[str]
+    torch_version: Optional[str]
+    torch_short: Optional[str]
+    torch_mm: Optional[str]
+    gpu_name: Optional[str] = None
+    gpu_compute: Optional[str] = None
+
+    @classmethod
+    def detect(cls, torch_version: Optional[str] = None) -> "RuntimeEnv":
+        """Detect runtime environment from current system."""
+        py_version = f"{sys.version_info.major}.{sys.version_info.minor}"
+        cuda_version = detect_cuda_version()
+        torch_version = torch_version or _detect_torch_version()
+
+        gpu_name, gpu_compute = None, None
+        try:
+            env = detect_cuda_environment()
+            if env.gpus:
+                gpu_name, gpu_compute = env.gpus[0].name, env.gpus[0].sm_version()
+        except Exception:
+            pass
+
+        return cls(
+            os_name=_get_os_name(),
+            platform_tag=get_platform_tag(),
+            python_version=py_version,
+            python_short=py_version.replace(".", ""),
+            cuda_version=cuda_version,
+            cuda_short=cuda_version.replace(".", "") if cuda_version else None,
+            torch_version=torch_version,
+            torch_short=torch_version.replace(".", "") if torch_version else None,
+            torch_mm="".join(torch_version.split(".")[:2]) if torch_version else None,
+            gpu_name=gpu_name,
+            gpu_compute=gpu_compute,
+        )
+
+    def as_dict(self) -> Dict[str, str]:
+        """Convert to dict for template substitution."""
+        result = {
+            "os": self.os_name,
+            "platform": self.platform_tag,
+            "python_version": self.python_version,
+            "py_version": self.python_version,
+            "py_short": self.python_short,
+            "py_minor": self.python_version.split(".")[-1],
+            "py_tag": f"cp{self.python_short}",
+        }
+        if self.cuda_version:
+            result.update(cuda_version=self.cuda_version, cuda_short=self.cuda_short,
+                         cuda_major=self.cuda_version.split(".")[0])
+        if self.torch_version:
+            result.update(torch_version=self.torch_version, torch_short=self.torch_short,
+                         torch_mm=self.torch_mm, torch_dotted_mm=".".join(self.torch_version.split(".")[:2]))
+        return result
+
+    def __str__(self) -> str:
+        parts = [f"Python {self.python_version}",
+                 f"CUDA {self.cuda_version}" if self.cuda_version else "CPU"]
+        if self.torch_version: parts.append(f"PyTorch {self.torch_version}")
+        if self.gpu_name: parts.append(f"GPU: {self.gpu_name}")
+        return ", ".join(parts)
+
+
+def detect_runtime(torch_version: Optional[str] = None) -> RuntimeEnv:
+    return RuntimeEnv.detect(torch_version)
+
+
+def _detect_torch_version() -> Optional[str]:
+    try:
+        import torch
+        return torch.__version__.split('+')[0]
+    except ImportError:
+        return None
+
+
+def parse_wheel_requirement(req: str) -> tuple[str, Optional[str]]:
+    """Parse 'pkg==1.0' -> ('pkg', '1.0')"""
+    for op in ['==', '>=', '<=', '~=', '!=', '>', '<']:
+        if op in req:
+            parts = req.split(op, 1)
+            return (parts[0].strip(), parts[1].strip())
+    return (req.strip(), None)

comfy_env/environment/__init__.py

@@ -0,0 +1,53 @@
+"""
+Environment layer - Environment management with side effects.
+
+Handles environment caching, path resolution, and runtime setup.
+"""
+
+from .cache import (
+    CACHE_DIR,
+    MARKER_FILE,
+    get_cache_dir,
+    get_env_name,
+    get_env_path,
+    cleanup_orphaned_envs,
+    write_marker_file,
+    read_marker_file,
+)
+from .paths import (
+    resolve_env_path,
+    get_site_packages_path,
+    get_lib_path,
+    copy_files,
+)
+from .setup import (
+    setup_env,
+    load_env_vars,
+    inject_site_packages,
+)
+from .libomp import (
+    dedupe_libomp,
+)
+
+__all__ = [
+    # Cache management
+    "CACHE_DIR",
+    "MARKER_FILE",
+    "get_cache_dir",
+    "get_env_name",
+    "get_env_path",
+    "cleanup_orphaned_envs",
+    "write_marker_file",
+    "read_marker_file",
+    # Path resolution
+    "resolve_env_path",
+    "get_site_packages_path",
+    "get_lib_path",
+    "copy_files",
+    # Setup helpers
+    "setup_env",
+    "load_env_vars",
+    "inject_site_packages",
+    # macOS workaround
+    "dedupe_libomp",
+]

comfy_env/environment/cache.py

@@ -0,0 +1,141 @@
+"""Central environment cache at ~/.comfy-env/envs/"""
+
+import glob
+import hashlib
+import shutil
+import sys
+from datetime import datetime
+from pathlib import Path
+from typing import Callable, Optional, Tuple
+
+import tomli
+import tomli_w
+
+try:
+    from .. import __version__
+except ImportError:
+    __version__ = "0.0.0-dev"
+
+CACHE_DIR = Path.home() / ".comfy-env" / "envs"
+MARKER_FILE = ".comfy-env-marker.toml"
+METADATA_FILE = ".comfy-env-metadata.toml"
+
+
+def get_cache_dir() -> Path:
+    CACHE_DIR.mkdir(parents=True, exist_ok=True)
+    return CACHE_DIR
+
+
+def compute_config_hash(config_path: Path) -> str:
+    return hashlib.sha256(config_path.read_bytes()).hexdigest()[:8]
+
+
+def sanitize_name(name: str) -> str:
+    name = name.lower()
+    for prefix in ("comfyui-", "comfyui_"):
+        if name.startswith(prefix): name = name[len(prefix):]
+    return name.replace("-", "_").replace(" ", "_")
+
+
+def get_env_name(node_dir: Path, config_path: Path) -> str:
+    """Generate env name: <nodename>_<subfolder>_<hash>"""
+    node_name = sanitize_name(node_dir.name)
+    config_parent = config_path.parent
+    if config_parent == node_dir:
+        subfolder = ""
+    else:
+        try:
+            subfolder = config_parent.relative_to(node_dir).as_posix().replace("/", "_")
+        except ValueError:
+            subfolder = sanitize_name(config_parent.name)
+    return f"{node_name}_{subfolder}_{compute_config_hash(config_path)}"
+
+
+def get_env_path(node_dir: Path, config_path: Path) -> Path:
+    return get_cache_dir() / get_env_name(node_dir, config_path)
+
+get_central_env_path = get_env_path
+
+
+def write_marker_file(config_path: Path, env_path: Path) -> None:
+    marker_path = config_path.parent / MARKER_FILE
+    marker_path.write_text(tomli_w.dumps({
+        "env": {"name": env_path.name, "path": str(env_path),
+                "config_hash": compute_config_hash(config_path),
+                "created": datetime.now().isoformat(), "comfy_env_version": __version__}
+    }))
+
+write_marker = write_marker_file
+
+
+def write_env_metadata(env_path: Path, marker_path: Path) -> None:
+    (env_path / METADATA_FILE).write_text(tomli_w.dumps({
+        "marker_path": str(marker_path), "created": datetime.now().isoformat()
+    }))
+
+
+def read_marker_file(marker_path: Path) -> Optional[dict]:
+    if not marker_path.exists(): return None
+    try:
+        with open(marker_path, "rb") as f: return tomli.load(f)
+    except Exception: return None
+
+read_marker = read_marker_file
+
+
+def read_env_metadata(env_path: Path) -> Optional[dict]:
+    metadata_path = env_path / METADATA_FILE
+    if not metadata_path.exists(): return None
+    try:
+        with open(metadata_path, "rb") as f: return tomli.load(f)
+    except Exception: return None
+
+
+def resolve_env_path(node_dir: Path) -> Tuple[Optional[Path], Optional[Path], Optional[Path]]:
+    """Resolve env with fallback: marker -> _env_<name> -> .pixi -> .venv"""
+    # Marker
+    marker = read_marker_file(node_dir / MARKER_FILE)
+    if marker and "env" in marker:
+        env_path = Path(marker["env"]["path"])
+        if env_path.exists(): return _get_env_paths(env_path)
+
+    # Local _env_<name>
+    local_env = node_dir / f"_env_{sanitize_name(node_dir.name)}"
+    if local_env.exists(): return _get_env_paths(local_env)
+
+    # .pixi
+    pixi_env = node_dir / ".pixi" / "envs" / "default"
+    if pixi_env.exists(): return _get_env_paths(pixi_env)
+
+    # .venv
+    venv = node_dir / ".venv"
+    if venv.exists(): return _get_env_paths(venv)
+
+    return None, None, None
+
+
+def _get_env_paths(env_path: Path) -> Tuple[Path, Optional[Path], Optional[Path]]:
+    if sys.platform == "win32":
+        return env_path, env_path / "Lib" / "site-packages", env_path / "Library" / "bin"
+    matches = glob.glob(str(env_path / "lib" / "python*" / "site-packages"))
+    return env_path, Path(matches[0]) if matches else None, env_path / "lib"
+
+
+def cleanup_orphaned_envs(log: Callable[[str], None] = print) -> int:
+    """Remove envs whose marker files no longer exist."""
+    cache_dir = get_cache_dir()
+    if not cache_dir.exists(): return 0
+
+    cleaned = 0
+    for env_dir in cache_dir.iterdir():
+        if not env_dir.is_dir(): continue
+        metadata = read_env_metadata(env_dir)
+        if not metadata: continue
+        marker_path = metadata.get("marker_path", "")
+        if marker_path and not Path(marker_path).exists():
+            log(f"[comfy-env] Cleaning: {env_dir.name}")
+            try:
+                shutil.rmtree(env_dir)
+                cleaned += 1
+            except Exception: pass
+    return cleaned

comfy_env/environment/libomp.py

@@ -0,0 +1,41 @@
+"""macOS: Dedupe libomp.dylib to prevent OpenMP runtime conflicts."""
+
+import glob
+import os
+import sys
+
+
+def dedupe_libomp() -> None:
+    """Symlink all libomp copies to torch's to prevent dual-runtime crashes."""
+    if sys.platform != "darwin":
+        return
+
+    try:
+        import torch
+        torch_libomp = os.path.join(os.path.dirname(torch.__file__), 'lib', 'libomp.dylib')
+        if not os.path.exists(torch_libomp):
+            return
+    except ImportError:
+        return
+
+    site_packages = os.path.dirname(os.path.dirname(torch.__file__))
+    patterns = [
+        os.path.join(site_packages, '*', 'Frameworks', 'libomp.dylib'),
+        os.path.join(site_packages, '*', '.dylibs', 'libomp.dylib'),
+        os.path.join(site_packages, '*', 'lib', 'libomp.dylib'),
+    ]
+
+    for pattern in patterns:
+        for libomp in glob.glob(pattern):
+            if 'torch' in libomp:
+                continue
+            if os.path.islink(libomp) and os.path.realpath(libomp) == os.path.realpath(torch_libomp):
+                continue
+            try:
+                if os.path.islink(libomp):
+                    os.unlink(libomp)
+                else:
+                    os.rename(libomp, libomp + '.bak')
+                os.symlink(torch_libomp, libomp)
+            except OSError:
+                pass

comfy_env/environment/paths.py

@@ -0,0 +1,38 @@
+"""Environment path utilities."""
+
+import shutil
+from pathlib import Path
+from typing import Optional, Tuple
+
+from .cache import resolve_env_path as _resolve_env_path
+
+
+def get_site_packages_path(node_dir: Path) -> Optional[Path]:
+    _, site_packages, _ = _resolve_env_path(node_dir)
+    return site_packages
+
+
+def get_lib_path(node_dir: Path) -> Optional[Path]:
+    _, _, lib_dir = _resolve_env_path(node_dir)
+    return lib_dir
+
+
+def resolve_env_path(node_dir: Path) -> Tuple[Optional[Path], Optional[Path], Optional[Path]]:
+    return _resolve_env_path(node_dir)
+
+
+def copy_files(src, dst, pattern: str = "*", overwrite: bool = False) -> int:
+    """Copy files matching pattern from src to dst."""
+    src, dst = Path(src), Path(dst)
+    if not src.exists(): return 0
+
+    dst.mkdir(parents=True, exist_ok=True)
+    copied = 0
+    for f in src.glob(pattern):
+        if f.is_file():
+            target = dst / f.relative_to(src)
+            target.parent.mkdir(parents=True, exist_ok=True)
+            if overwrite or not target.exists():
+                shutil.copy2(f, target)
+                copied += 1
+    return copied

comfy_env/environment/setup.py

@@ -0,0 +1,88 @@
+"""Environment setup for ComfyUI prestartup."""
+
+import glob
+import os
+import sys
+from pathlib import Path
+from typing import Dict, Optional
+
+from .cache import MARKER_FILE, sanitize_name
+from .libomp import dedupe_libomp
+
+USE_COMFY_ENV_VAR = "USE_COMFY_ENV"
+
+
+def is_comfy_env_enabled() -> bool:
+    return os.environ.get(USE_COMFY_ENV_VAR, "1").lower() not in ("0", "false", "no", "off")
+
+
+def load_env_vars(config_path: str) -> Dict[str, str]:
+    """Load [env_vars] from comfy-env.toml."""
+    if not os.path.exists(config_path): return {}
+    try:
+        import tomli
+        with open(config_path, "rb") as f:
+            return {str(k): str(v) for k, v in tomli.load(f).get("env_vars", {}).items()}
+    except Exception:
+        return {}
+
+
+def inject_site_packages(env_path: str) -> Optional[str]:
+    """Add site-packages to sys.path."""
+    if sys.platform == "win32":
+        site_packages = os.path.join(env_path, "Lib", "site-packages")
+    else:
+        matches = glob.glob(os.path.join(env_path, "lib", "python*", "site-packages"))
+        site_packages = matches[0] if matches else None
+
+    if site_packages and os.path.exists(site_packages) and site_packages not in sys.path:
+        sys.path.insert(0, site_packages)
+        return site_packages
+    return None
+
+
+def setup_env(node_dir: Optional[str] = None) -> None:
+    """Set up env for pixi libraries. Call in prestartup_script.py before native imports."""
+    if not is_comfy_env_enabled(): return
+    dedupe_libomp()
+
+    if node_dir is None:
+        import inspect
+        node_dir = str(Path(inspect.stack()[1].filename).parent)
+
+    # Apply env vars
+    for k, v in load_env_vars(os.path.join(node_dir, "comfy-env.toml")).items():
+        os.environ[k] = v
+
+    # Find env: marker -> _env_<name> -> .pixi
+    pixi_env = None
+    marker_path = os.path.join(node_dir, MARKER_FILE)
+    if os.path.exists(marker_path):
+        try:
+            import tomli
+            with open(marker_path, "rb") as f:
+                env_path = tomli.load(f).get("env", {}).get("path")
+            if env_path and os.path.exists(env_path):
+                pixi_env = env_path
+        except Exception: pass
+
+    if not pixi_env:
+        local_env = os.path.join(node_dir, f"_env_{sanitize_name(os.path.basename(node_dir))}")
+        if os.path.exists(local_env): pixi_env = local_env
+
+    if not pixi_env:
+        old_pixi = os.path.join(node_dir, ".pixi", "envs", "default")
+        if os.path.exists(old_pixi): pixi_env = old_pixi
+
+    if not pixi_env: return
+
+    # Set library paths
+    if sys.platform == "win32":
+        lib_dir = os.path.join(pixi_env, "Library", "bin")
+        if os.path.exists(lib_dir): os.environ["PATH"] = lib_dir + ";" + os.environ.get("PATH", "")
+    else:
+        lib_dir = os.path.join(pixi_env, "lib")
+        var = "DYLD_LIBRARY_PATH" if sys.platform == "darwin" else "LD_LIBRARY_PATH"
+        if os.path.exists(lib_dir): os.environ[var] = lib_dir + ":" + os.environ.get(var, "")
+
+    inject_site_packages(pixi_env)