comfy-env 0.0.73__tar.gz → 0.0.75__tar.gz

{comfy_env-0.0.73 → comfy_env-0.0.75}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: comfy-env
-Version: 0.0.73
+Version: 0.0.75
 Summary: Environment management for ComfyUI custom nodes - CUDA wheel resolution and process isolation
 Project-URL: Homepage, https://github.com/PozzettiAndrea/comfy-env
 Project-URL: Repository, https://github.com/PozzettiAndrea/comfy-env

{comfy_env-0.0.73 → comfy_env-0.0.75}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "comfy-env"
-version = "0.0.73"
+version = "0.0.75"
 description = "Environment management for ComfyUI custom nodes - CUDA wheel resolution and process isolation"
 readme = "README.md"
 license = {text = "MIT"}

{comfy_env-0.0.73 → comfy_env-0.0.75}/src/comfy_env/install.py

@@ -59,6 +59,10 @@ def install(
     if cfg.apt_packages:
         _install_apt_packages(cfg.apt_packages, log, dry_run)
 
+    # Set persistent env vars (for OpenMP settings, etc.)
+    if cfg.env_vars:
+        _set_persistent_env_vars(cfg.env_vars, log, dry_run)
+
     # Install node dependencies
     if cfg.node_reqs:
         _install_node_dependencies(cfg.node_reqs, node_dir, log, dry_run)
@@ -117,20 +121,116 @@ def _install_apt_packages(
     if result.returncode != 0:
         log(f"[apt] Warning: apt-get update failed: {result.stderr.strip()}")
 
-    # Run apt-get install
-    log(f"[apt] Installing packages...")
-    result = subprocess.run(
-        prefix + ["apt-get", "install", "-y"] + packages,
-        capture_output=True,
-        text=True,
-    )
+    # 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
+
+    if not env_vars:
+        return
+
+    system = platform.system()
+    log(f"\n[env] Setting {len(env_vars)} persistent environment variable(s)...")
 
-    if result.returncode == 0:
-        log("[apt] System packages installed successfully!")
+    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)
+        _add_to_shell_profile(env_vars, log)
+
+    else:  # Linux
+        _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
+    shell = os.environ.get("SHELL", "/bin/bash")
+    if "zsh" in shell:
+        rc_file = Path.home() / ".zshrc"
+    else:
+        rc_file = Path.home() / ".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}")
+
+    # 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:
+        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"[apt] Warning: Installation failed: {result.stderr.strip()}")
-        log(f"[apt] You may need to install manually:")
-        log(f"  sudo apt-get install -y {' '.join(packages)}")
+        log(f"  [env] Already configured in {rc_file}")
+
+    log("[env] Restart terminal/ComfyUI for changes to take effect")
 
 
 def _install_node_dependencies(

{comfy_env-0.0.73 → comfy_env-0.0.75}/src/comfy_env/pixi/core.py

@@ -31,6 +31,12 @@ PIXI_URLS = {
 # CUDA wheels index
 CUDA_WHEELS_INDEX = "https://pozzettiandrea.github.io/cuda-wheels/"
 
+# Flash attention pre-built wheels (mjun0812)
+FLASH_ATTN_INDEX = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/expanded_assets/v0.7.16"
+
+# PyTorch Geometric wheels index
+PYG_WHEELS_INDEX = "https://data.pyg.org/whl/"
+
 # CUDA version -> PyTorch version mapping
 CUDA_TORCH_MAP = {
     "12.8": "2.8",
@@ -39,6 +45,20 @@ CUDA_TORCH_MAP = {
 }
 
 
+def get_all_find_links(package: str, torch_version: str, cuda_version: str) -> list:
+    """Get all find-links URLs for a CUDA package."""
+    cuda_short = cuda_version.replace(".", "")[:3]
+
+    return [
+        # cuda-wheels index (package-specific page)
+        f"{CUDA_WHEELS_INDEX}{package}/",
+        # flash-attn GitHub releases
+        FLASH_ATTN_INDEX,
+        # PyG index (torch version specific)
+        f"{PYG_WHEELS_INDEX}torch-{torch_version}.0+cu{cuda_short}.html",
+    ]
+
+
 def get_current_platform() -> str:
     """Get the current platform string for pixi."""
     if sys.platform == "linux":
@@ -310,7 +330,7 @@ def pixi_install(
         log(f"pixi install failed:\n{result.stderr}")
         raise RuntimeError(f"pixi install failed: {result.stderr}")
 
-    # Install CUDA packages with --no-index --find-links (bypasses PyPI completely)
+    # Install CUDA packages with --find-links (searches all known sources)
     if cfg.cuda_packages and cuda_version:
         log(f"Installing CUDA packages: {cfg.cuda_packages}")
         python_path = get_pixi_python(node_dir)
@@ -318,18 +338,18 @@ def pixi_install(
             raise RuntimeError("Could not find Python in pixi environment")
 
         for package in cfg.cuda_packages:
-            # Each package has its own find-links page at CUDA_WHEELS_INDEX/<package>/
-            find_links_url = f"{CUDA_WHEELS_INDEX}{package}/"
-            log(f"  Installing {package} from {find_links_url}")
-
             pip_cmd = [
                 str(python_path), "-m", "pip", "install",
-                "--no-index",
                 "--no-deps",
                 "--no-cache-dir",
-                "--find-links", find_links_url,
-                package,
             ]
+
+            # Add all find-links sources
+            for url in get_all_find_links(package, torch_version, cuda_version):
+                pip_cmd.extend(["--find-links", url])
+
+            log(f"  Installing {package}")
+            pip_cmd.append(package)
             result = subprocess.run(pip_cmd, capture_output=True, text=True)
             if result.returncode != 0:
                 log(f"CUDA package install failed for {package}:\n{result.stderr}")