comfy-env 0.0.12__py3-none-any.whl → 0.0.13__py3-none-any.whl

comfy_env/cli.py

@@ -234,6 +234,7 @@ def cmd_info(args) -> int:
 def cmd_resolve(args) -> int:
     """Handle resolve command."""
     from .resolver import RuntimeEnv, WheelResolver, parse_wheel_requirement
+    from .registry import PACKAGE_REGISTRY
     from .env.config_file import discover_env_config, load_env_from_file
 
     env = RuntimeEnv.detect()
@@ -270,19 +271,116 @@ def cmd_resolve(args) -> int:
             print(f"  {package}: No version specified, skipping")
             continue
 
+        pkg_lower = package.lower()
         try:
-            url = resolver.resolve(package, version, env, verify=args.verify)
-            status = "OK" if args.verify else "resolved"
-            print(f"  {package}=={version}: {status}")
-            print(f"    {url}")
+            # Check if package is in registry with github_release method
+            if pkg_lower in PACKAGE_REGISTRY:
+                registry_config = PACKAGE_REGISTRY[pkg_lower]
+                method = registry_config.get("method")
+
+                if method == "github_release":
+                    # Resolve URL from registry sources
+                    url = _resolve_github_release_url(package, version, env, registry_config)
+                    status = "OK" if args.verify else "resolved"
+                    print(f"  {package}=={version}: {status}")
+                    print(f"    {url}")
+                else:
+                    # For other methods, just show what method will be used
+                    print(f"  {package}=={version}: uses {method} method")
+                    if "index_url" in registry_config:
+                        index_url = _substitute_template(registry_config["index_url"], env)
+                        print(f"    index: {index_url}")
+                    elif "package_template" in registry_config:
+                        pkg_name = _substitute_template(registry_config["package_template"], env)
+                        print(f"    installs as: {pkg_name}")
+            else:
+                # Fall back to WheelResolver
+                url = resolver.resolve(package, version, env, verify=args.verify)
+                status = "OK" if args.verify else "resolved"
+                print(f"  {package}=={version}: {status}")
+                print(f"    {url}")
         except Exception as e:
             print(f"  {package}=={version}: FAILED")
-            print(f"    {e}")
+            _print_wheel_not_found_error(package, version, env, e)
             all_ok = False
 
     return 0 if all_ok else 1
 
 
+def _substitute_template(template: str, env) -> str:
+    """Substitute environment variables into a URL template."""
+    vars_dict = env.as_dict()
+    result = template
+    for key, value in vars_dict.items():
+        if value is not None:
+            result = result.replace(f"{{{key}}}", str(value))
+    return result
+
+
+def _resolve_github_release_url(package: str, version: str, env, config: dict) -> str:
+    """Resolve URL for github_release method packages."""
+    sources = config.get("sources", [])
+    if not sources:
+        raise ValueError(f"No sources configured for {package}")
+
+    # Build template variables
+    vars_dict = env.as_dict()
+    vars_dict["version"] = version
+    vars_dict["py_tag"] = f"cp{env.python_short}"
+    if env.cuda_version:
+        vars_dict["cuda_major"] = env.cuda_version.split(".")[0]
+
+    # Filter sources by platform
+    current_platform = env.platform_tag
+    compatible_sources = [
+        s for s in sources
+        if current_platform in s.get("platforms", [])
+    ]
+
+    if not compatible_sources:
+        available = set()
+        for s in sources:
+            available.update(s.get("platforms", []))
+        raise ValueError(
+            f"No {package} wheels for platform {current_platform}. "
+            f"Available: {', '.join(sorted(available))}"
+        )
+
+    # Return URL from first compatible source
+    source = compatible_sources[0]
+    url_template = source.get("url_template", "")
+    url = url_template
+    for key, value in vars_dict.items():
+        if value is not None:
+            url = url.replace(f"{{{key}}}", str(value))
+
+    return url
+
+
+def _print_wheel_not_found_error(package: str, version: str, env, error: Exception) -> None:
+    """Print a formatted error message for wheel not found."""
+    from .errors import WheelNotFoundError
+
+    if isinstance(error, WheelNotFoundError):
+        print(f"    CUDA wheel not found: {package}=={version}")
+        print()
+        print("+------------------------------------------------------------------+")
+        print("|  CUDA Wheel Not Found                                            |")
+        print("+------------------------------------------------------------------+")
+        print(f"|  Package:   {package}=={version:<46} |")
+        print(f"|  Requested: cu{env.cuda_short}-torch{env.torch_mm}-{env.python_short}-{env.platform_tag:<17} |")
+        print("|                                                                  |")
+        print(f"|  Reason: {error.reason:<54} |")
+        print("|                                                                  |")
+        print("|  Suggestions:                                                    |")
+        print(f"|    1. Check if wheel exists: comfy-env resolve {package:<15} |")
+        print(f"|    2. Build wheel locally: comfy-env build {package:<18} |")
+        print("|                                                                  |")
+        print("+------------------------------------------------------------------+")
+    else:
+        print(f"    {error}")
+
+
 def cmd_doctor(args) -> int:
     """Handle doctor command."""
     from .install import verify_installation

comfy_env/env/config_file.py

@@ -168,7 +168,7 @@ def _get_default_pytorch_version(cuda_version: Optional[str]) -> str:
         - CUDA 12.8 (Turing+): PyTorch 2.8.0
     """
     if cuda_version == "12.4":
-        return "2.5.1"  # Legacy: Pascal GPUs
+        return "2.5.1"  # Full: Pascal GPUs
     return "2.8.0"  # Modern: Turing through Blackwell
 
 
@@ -176,7 +176,7 @@ def _parse_config(data: Dict[str, Any], base_dir: Path) -> IsolatedEnv:
     """
     Parse TOML data into IsolatedEnv.
 
-    Supports both simplified and legacy config formats:
+    Supports both simplified and full config formats:
 
     Simplified (CUDA packages only):
         [packages]
@@ -186,7 +186,7 @@ def _parse_config(data: Dict[str, Any], base_dir: Path) -> IsolatedEnv:
     Or as list:
         packages = ["torch-scatter==2.1.2", "torch-cluster==1.6.3"]
 
-    Legacy:
+    Full:
         [env]
         name = "my-node"
         [packages]
@@ -233,7 +233,7 @@ def _parse_config(data: Dict[str, Any], base_dir: Path) -> IsolatedEnv:
         variables.setdefault("pytorch_mm", pytorch_mm)
 
     # Parse CUDA packages - support multiple formats
-    # Priority: [cuda] section > cuda = [...] > legacy [packages] section
+    # Priority: [cuda] section > cuda = [...] > [packages] section
     no_deps_requirements = []
     requirements = []
 
@@ -258,12 +258,12 @@ def _parse_config(data: Dict[str, Any], base_dir: Path) -> IsolatedEnv:
 
     elif isinstance(packages_section, dict):
         # Check for simplified format: [packages] with key=value pairs
-        # vs legacy format: [packages] with requirements/no_deps lists
+        # vs old format: [packages] with requirements/no_deps lists
 
-        has_legacy_keys = any(k in packages_section for k in ["requirements", "no_deps", "requirements_file"])
+        has_old_keys = any(k in packages_section for k in ["requirements", "no_deps", "requirements_file"])
 
-        if has_legacy_keys:
-            # Legacy format
+        if has_old_keys:
+            # Old format
             raw_requirements = packages_section.get("requirements", [])
             requirements = [_substitute_vars(req, variables) for req in raw_requirements]
 
@@ -340,9 +340,9 @@ def load_config(
     base_dir: Optional[Path] = None,
 ) -> EnvManagerConfig:
     """
-    Load full EnvManagerConfig from a TOML file (v2 schema).
+    Load full EnvManagerConfig from a TOML file.
 
-    Supports both v2 schema and legacy format (auto-detected).
+    Supports both full schema (named envs) and simple format (auto-detected).
 
     Args:
         path: Path to the TOML config file
@@ -373,7 +373,7 @@ def load_config(
     with open(path, "rb") as f:
         data = tomllib.load(f)
 
-    return _parse_config_v2(data, base_dir)
+    return _parse_full_config(data, base_dir)
 
 
 def discover_config(
@@ -404,9 +404,9 @@ def discover_config(
     return None
 
 
-def _parse_config_v2(data: Dict[str, Any], base_dir: Path) -> EnvManagerConfig:
+def _parse_full_config(data: Dict[str, Any], base_dir: Path) -> EnvManagerConfig:
     """
-    Parse TOML data into EnvManagerConfig (v2 schema).
+    Parse TOML data into EnvManagerConfig.
 
     Schema:
         [local.cuda]        - CUDA packages for host
@@ -416,16 +416,16 @@ def _parse_config_v2(data: Dict[str, Any], base_dir: Path) -> EnvManagerConfig:
         [envname.packages]  - Packages for env
         [node_reqs]         - Node dependencies
 
-    Also supports legacy format for backward compatibility.
+    Also supports simple format ([env] + [packages]) for backward compatibility.
     """
-    # Detect if this is v2 schema or legacy
-    is_v2 = "local" in data or _has_named_env(data)
+    # Detect if this is full schema or simple format
+    is_full = "local" in data or _has_named_env(data)
 
-    if not is_v2:
-        # Legacy format - convert to v2 structure
-        return _convert_legacy_to_v2(data, base_dir)
+    if not is_full:
+        # Simple format - convert to full structure
+        return _convert_simple_to_full(data, base_dir)
 
-    # Parse v2 schema
+    # Parse full schema
     local = _parse_local_section(data.get("local", {}))
     envs = _parse_env_sections(data, base_dir)
     node_reqs = _parse_node_reqs(data.get("node_reqs", {}))
@@ -582,27 +582,30 @@ def _parse_node_reqs(node_reqs_data: Dict[str, Any]) -> List[NodeReq]:
     return reqs
 
 
-def _convert_legacy_to_v2(data: Dict[str, Any], base_dir: Path) -> EnvManagerConfig:
-    """Convert legacy config format to v2 EnvManagerConfig."""
-    # Parse using legacy parser to get IsolatedEnv
-    legacy_env = _parse_config(data, base_dir)
+def _convert_simple_to_full(data: Dict[str, Any], base_dir: Path) -> EnvManagerConfig:
+    """Convert simple config format to full EnvManagerConfig.
 
-    # Check if this is really just a Type 2 config (local CUDA only, no venv)
-    # Type 2 indicators: no [env] section with name, or name matches directory
+    Simple configs have [env] and [packages] sections but no named environments.
+    This converts them to the full format with a single named environment.
+    """
+    # Parse using simple parser to get IsolatedEnv
+    simple_env = _parse_config(data, base_dir)
+
+    # Check if this has explicit env settings (isolated venv) vs just CUDA packages (local install)
     env_section = data.get("env", {})
     has_explicit_env = bool(env_section.get("name") or env_section.get("python"))
 
     if has_explicit_env:
-        # This is a Type 1 config (isolated venv)
+        # Isolated venv config
         return EnvManagerConfig(
             local=LocalConfig(),
-            envs={legacy_env.name: legacy_env},
+            envs={simple_env.name: simple_env},
             node_reqs=[],
         )
     else:
-        # This is a Type 2 config (local CUDA only)
+        # Local CUDA packages only (no isolated venv)
         cuda_packages = {}
-        for req in legacy_env.no_deps_requirements:
+        for req in simple_env.no_deps_requirements:
             if "==" in req:
                 pkg, ver = req.split("==", 1)
                 cuda_packages[pkg] = ver
@@ -612,8 +615,8 @@ def _convert_legacy_to_v2(data: Dict[str, Any], base_dir: Path) -> EnvManagerCon
         return EnvManagerConfig(
             local=LocalConfig(
                 cuda_packages=cuda_packages,
-                requirements=legacy_env.requirements,
+                requirements=simple_env.requirements,
             ),
             envs={},
             node_reqs=[],
-        )
+        )

comfy_env/env/detection.py

@@ -92,22 +92,6 @@ def is_blackwell_gpu(name: str, compute_cap: str) -> bool:
     return False
 
 
-def needs_cuda_128() -> bool:
-    """
-    Check if any detected GPU requires CUDA 12.8.
-
-    Returns:
-        True if Blackwell GPU detected, False otherwise.
-    """
-    gpus = detect_gpu_info()
-
-    for gpu in gpus:
-        if is_blackwell_gpu(gpu["name"], gpu["compute_cap"]):
-            return True
-
-    return False
-
-
 def is_legacy_gpu(compute_cap: str) -> bool:
     """
     Check if GPU is Pascal or older (requires legacy CUDA/PyTorch).

comfy_env/errors.py

@@ -291,35 +291,3 @@ class InstallError(EnvManagerError):
 
         details = "\n".join(details_parts) if details_parts else None
         super().__init__(message, details)
-
-
-def format_environment_mismatch(
-    expected: "RuntimeEnv",
-    actual: "RuntimeEnv",
-) -> str:
-    """
-    Format a message explaining environment mismatch.
-
-    Used when the current environment doesn't match what's needed.
-    """
-    mismatches = []
-
-    if expected.cuda_version != actual.cuda_version:
-        mismatches.append(
-            f"  CUDA: expected {expected.cuda_version}, got {actual.cuda_version}"
-        )
-
-    if expected.torch_version != actual.torch_version:
-        mismatches.append(
-            f"  PyTorch: expected {expected.torch_version}, got {actual.torch_version}"
-        )
-
-    if expected.python_version != actual.python_version:
-        mismatches.append(
-            f"  Python: expected {expected.python_version}, got {actual.python_version}"
-        )
-
-    if not mismatches:
-        return "Environment matches expected configuration"
-
-    return "Environment mismatch:\n" + "\n".join(mismatches)

comfy_env/install.py

@@ -421,6 +421,7 @@ def _install_from_github_release(
                 url = url.replace(f"{{{key}}}", str(value))
 
         log(f"  Trying {source_name}: {package}=={version}...")
+        log(f"    Resolved wheel to: {url}")
 
         try:
             pip_cmd = _get_pip_command()

comfy_env/ipc/bridge.py

@@ -48,10 +48,6 @@ class WorkerBridge:
         result = bridge.call("process", image=my_image)
     """
 
-    # Singleton instances by environment hash
-    _instances: Dict[str, "WorkerBridge"] = {}
-    _instances_lock = threading.Lock()
-
     def __init__(
         self,
         env: IsolatedEnv,
@@ -81,38 +77,6 @@ class WorkerBridge:
         self._process_lock = threading.Lock()
         self._stderr_thread: Optional[threading.Thread] = None
 
-    @classmethod
-    def get_instance(
-        cls,
-        env: IsolatedEnv,
-        worker_script: Path,
-        base_dir: Optional[Path] = None,
-        log_callback: Optional[Callable[[str], None]] = None,
-    ) -> "WorkerBridge":
-        """
-        Get or create a singleton bridge instance for an environment.
-
-        Args:
-            env: Isolated environment configuration
-            worker_script: Path to the worker Python script
-            base_dir: Base directory for environments
-            log_callback: Optional callback for logging
-
-        Returns:
-            WorkerBridge instance (reused if same env hash)
-        """
-        env_hash = env.get_env_hash()
-
-        with cls._instances_lock:
-            if env_hash not in cls._instances:
-                cls._instances[env_hash] = cls(
-                    env=env,
-                    worker_script=worker_script,
-                    base_dir=base_dir,
-                    log_callback=log_callback,
-                )
-            return cls._instances[env_hash]
-
     @classmethod
     def from_config_file(
         cls,

comfy_env/registry.py

@@ -1,8 +1,7 @@
 """Built-in registry of CUDA packages and their wheel sources.
 
-This module provides a mapping of well-known CUDA packages to their
-installation sources, eliminating the need for users to specify
-wheel_sources in their comfyui_env.toml.
+This module loads package configurations from wheel_sources.yml and provides
+lookup functions for the install module.
 
 Install method types:
 - "index": Use pip --extra-index-url (PEP 503 simple repository)
@@ -10,9 +9,13 @@ Install method types:
 - "find_links": Use pip --find-links (for PyG, etc.)
 - "pypi_variant": Package name varies by CUDA version (e.g., spconv-cu124)
 - "github_release": Direct wheel URL from GitHub releases with fallback sources
+- "pypi": Standard PyPI install
 """
 
-from typing import Dict, Any, Optional
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import yaml
 
 
 def get_cuda_short2(cuda_version: str) -> str:
@@ -40,180 +43,16 @@ def get_cuda_short2(cuda_version: str) -> str:
     return f"{major}{minor}"
 
 
-# =============================================================================
-# Package Registry
-# =============================================================================
-# Maps package names to their installation configuration.
-#
-# Template variables available:
-#   {cuda_version}  - Full CUDA version (e.g., "12.8")
-#   {cuda_short}    - CUDA without dot (e.g., "128")
-#   {cuda_short2}   - CUDA short for spconv (e.g., "124" not "1240")
-#   {torch_version} - Full PyTorch version (e.g., "2.8.0")
-#   {torch_short}   - PyTorch without dots (e.g., "280")
-#   {torch_mm}      - PyTorch major.minor (e.g., "28")
-#   {py_version}    - Python version (e.g., "3.10")
-#   {py_short}      - Python without dot (e.g., "310")
-#   {py_minor}      - Python minor version only (e.g., "10")
-#   {platform}      - Platform tag (e.g., "linux_x86_64")
-# =============================================================================
-
-PACKAGE_REGISTRY: Dict[str, Dict[str, Any]] = {
-    # =========================================================================
-    # PyTorch Geometric (PyG) packages - official index
-    # https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html
-    # Uses --find-links (not --extra-index-url) for proper wheel discovery
-    # =========================================================================
-    "torch-scatter": {
-        "method": "find_links",
-        "index_url": "https://data.pyg.org/whl/torch-{torch_version}+cu{cuda_short}.html",
-        "description": "Scatter operations for PyTorch",
-    },
-    "torch-cluster": {
-        "method": "find_links",
-        "index_url": "https://data.pyg.org/whl/torch-{torch_version}+cu{cuda_short}.html",
-        "description": "Clustering algorithms for PyTorch",
-    },
-    "torch-sparse": {
-        "method": "find_links",
-        "index_url": "https://data.pyg.org/whl/torch-{torch_version}+cu{cuda_short}.html",
-        "description": "Sparse tensor operations for PyTorch",
-    },
-    "torch-spline-conv": {
-        "method": "find_links",
-        "index_url": "https://data.pyg.org/whl/torch-{torch_version}+cu{cuda_short}.html",
-        "description": "Spline convolutions for PyTorch",
-    },
-
-    # =========================================================================
-    # pytorch3d - Facebook's official wheels
-    # https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md
-    # =========================================================================
-    "pytorch3d": {
-        "method": "index",
-        "index_url": "https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py3{py_minor}_cu{cuda_short}_pyt{torch_short}/download.html",
-        "description": "PyTorch3D - 3D deep learning library",
-    },
-
-    # =========================================================================
-    # PozzettiAndrea wheel repos (GitHub Pages indexes)
-    # =========================================================================
-    # nvdiffrast - wheels are now at cu{cuda}-torch{torch_short} releases
-    "nvdiffrast": {
-        "method": "github_index",
-        "index_url": "https://pozzettiandrea.github.io/nvdiffrast-full-wheels/cu{cuda_short}-torch{torch_short}/",
-        "description": "NVIDIA differentiable rasterizer",
-    },
-    # cumesh, o_voxel, flex_gemm, nvdiffrec_render use torch_short (3 digits: 280)
-    "cumesh": {
-        "method": "github_index",
-        "index_url": "https://pozzettiandrea.github.io/cumesh-wheels/cu{cuda_short}-torch{torch_short}/",
-        "description": "CUDA-accelerated mesh utilities",
-    },
-    "o_voxel": {
-        "method": "github_index",
-        "index_url": "https://pozzettiandrea.github.io/ovoxel-wheels/cu{cuda_short}-torch{torch_short}/",
-        "description": "O-Voxel CUDA extension for TRELLIS",
-    },
-    "flex_gemm": {
-        "method": "github_index",
-        "index_url": "https://pozzettiandrea.github.io/flexgemm-wheels/cu{cuda_short}-torch{torch_short}/",
-        "description": "Flexible GEMM operations",
-    },
-    "nvdiffrec_render": {
-        "method": "github_release",
-        "sources": [
-            {
-                "name": "PozzettiAndrea",
-                "url_template": "https://github.com/PozzettiAndrea/nvdiffrec_render-wheels/releases/download/cu{cuda_short}-torch{torch_short}/nvdiffrec_render-{version}%2Bcu{cuda_short}torch{torch_mm}-{py_tag}-{py_tag}-linux_x86_64.whl",
-                "platforms": ["linux_x86_64"],
-            },
-            {
-                "name": "PozzettiAndrea-windows",
-                "url_template": "https://github.com/PozzettiAndrea/nvdiffrec_render-wheels/releases/download/cu{cuda_short}-torch{torch_short}/nvdiffrec_render-{version}%2Bcu{cuda_short}torch{torch_mm}-{py_tag}-{py_tag}-win_amd64.whl",
-                "platforms": ["win_amd64", "windows_amd64"],
-            },
-        ],
-        "description": "NVDiffRec rendering utilities",
-    },
-
-    # =========================================================================
-    # spconv - PyPI with CUDA-versioned package names
-    # Package names: spconv-cu118, spconv-cu120, spconv-cu121, spconv-cu124, spconv-cu126
-    # Note: Max available is cu126 as of Jan 2026, use explicit spconv-cu126 in config
-    # =========================================================================
-    "spconv": {
-        "method": "pypi_variant",
-        "package_template": "spconv-cu{cuda_short2}",
-        "description": "Sparse convolution library (use spconv-cu126 for CUDA 12.6+)",
-    },
-
-    # =========================================================================
-    # sageattention - Fast quantized attention (2-5x faster than FlashAttention)
-    # Linux: Prebuilt wheels from Kijai/PrecompiledWheels (v2.2.0, cp312)
-    # Windows: Prebuilt wheels from woct0rdho (v2.2.0, cp39-abi3)
-    # =========================================================================
-    "sageattention": {
-        "method": "github_release",
-        "sources": [
-            # Linux: Kijai's precompiled wheels on HuggingFace (Python 3.12)
-            {
-                "name": "kijai-hf",
-                "url_template": "https://huggingface.co/Kijai/PrecompiledWheels/resolve/main/sageattention-{version}-cp312-cp312-linux_x86_64.whl",
-                "platforms": ["linux_x86_64"],
-            },
-            # Windows: woct0rdho prebuilt wheels (ABI3: Python >= 3.9)
-            # Format: sageattention-2.2.0+cu128torch2.8.0.post3-cp39-abi3-win_amd64.whl
-            {
-                "name": "woct0rdho",
-                "url_template": "https://github.com/woct0rdho/SageAttention/releases/download/v2.2.0-windows.post3/sageattention-2.2.0%2Bcu{cuda_short}torch{torch_version}.post3-cp39-abi3-win_amd64.whl",
-                "platforms": ["win_amd64"],
-            },
-        ],
-        "description": "SageAttention - 2-5x faster than FlashAttention with quantized kernels",
-    },
-
-    # =========================================================================
-    # triton - Required for sageattention on Linux (usually bundled with PyTorch)
-    # =========================================================================
-    "triton": {
-        "method": "pypi",
-        "description": "Triton compiler for custom CUDA kernels (required by sageattention)",
-    },
-
-    # =========================================================================
-    # flash-attn - Multi-source prebuilt wheels
-    # Required for UniRig and other transformer-based models
-    # Sources: Dao-AILab (official), mjun0812 (Linux), bdashore3 (Windows)
-    # =========================================================================
-    "flash-attn": {
-        "method": "github_release",
-        "sources": [
-            # Linux: Dao-AILab official wheels (CUDA 12.x, PyTorch 2.4-2.8)
-            # Format: flash_attn-2.8.3+cu12torch2.8cxx11abiTRUE-cp310-cp310-linux_x86_64.whl
-            {
-                "name": "Dao-AILab",
-                "url_template": "https://github.com/Dao-AILab/flash-attention/releases/download/v{version}/flash_attn-{version}%2Bcu{cuda_major}torch{torch_dotted_mm}cxx11abiTRUE-{py_tag}-{py_tag}-linux_x86_64.whl",
-                "platforms": ["linux_x86_64"],
-            },
-            # Linux: mjun0812 prebuilt wheels (CUDA 12.4-13.0, PyTorch 2.5-2.9)
-            # Format: flash_attn-2.8.3+cu128torch2.8-cp310-cp310-linux_x86_64.whl
-            # Note: Release v0.7.2 contains multiple flash_attn versions
-            {
-                "name": "mjun0812",
-                "url_template": "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.7.2/flash_attn-{version}%2Bcu{cuda_short}torch{torch_dotted_mm}-{py_tag}-{py_tag}-linux_x86_64.whl",
-                "platforms": ["linux_x86_64"],
-            },
-            # Windows: bdashore3 prebuilt wheels (CUDA 12.4/12.8, PyTorch 2.6-2.8)
-            {
-                "name": "bdashore3",
-                "url_template": "https://github.com/bdashore3/flash-attention/releases/download/v{version}/flash_attn-{version}%2Bcu{cuda_short}torch{torch_version}cxx11abiFALSE-{py_tag}-{py_tag}-win_amd64.whl",
-                "platforms": ["win_amd64"],
-            },
-        ],
-        "description": "Flash Attention for fast transformer inference",
-    },
-}
+def _load_wheel_sources() -> Dict[str, Dict[str, Any]]:
+    """Load package registry from wheel_sources.yml."""
+    yml_path = Path(__file__).parent / "wheel_sources.yml"
+    with open(yml_path, "r") as f:
+        data = yaml.safe_load(f)
+    return data.get("packages", {})
+
+
+# Load registry at module import time
+PACKAGE_REGISTRY: Dict[str, Dict[str, Any]] = _load_wheel_sources()
 
 
 def get_package_info(package: str) -> Optional[Dict[str, Any]]:

comfy_env/resolver.py

@@ -218,21 +218,6 @@ class WheelSource:
         return package.lower() in [p.lower() for p in self.packages]
 
 
-# Default wheel sources for common CUDA packages
-DEFAULT_WHEEL_SOURCES = [
-    WheelSource(
-        name="nvdiffrast-wheels",
-        url_template="https://github.com/PozzettiAndrea/nvdiffrast-full-wheels/releases/download/v{version}/nvdiffrast-{version}%2Bcu{cuda_short}torch{torch_mm}-cp{py_short}-cp{py_short}-{platform}.whl",
-        packages=["nvdiffrast"],
-    ),
-    WheelSource(
-        name="cumesh-wheels",
-        url_template="https://github.com/PozzettiAndrea/cumesh-wheels/releases/download/v{version}/{package}-{version}%2Bcu{cuda_short}torch{torch_mm}-cp{py_short}-cp{py_short}-{platform}.whl",
-        packages=["pytorch3d", "torch-cluster", "torch-scatter", "torch-sparse"],
-    ),
-]
-
-
 class WheelResolver:
     """
     Resolves CUDA wheel URLs from package name and runtime environment.
@@ -255,10 +240,11 @@ class WheelResolver:
         Initialize resolver.
 
         Args:
-            sources: List of WheelSource configurations.
+            sources: List of WheelSource configurations. Defaults to empty
+                    (use PACKAGE_REGISTRY in install.py for actual sources).
             overrides: Package-specific URL overrides (package -> template).
         """
-        self.sources = sources or DEFAULT_WHEEL_SOURCES
+        self.sources = sources or []
         self.overrides = overrides or {}
 
     def resolve(

comfy_env/wheel_sources.yml

@@ -0,0 +1,141 @@
+# Wheel sources registry for CUDA packages
+#
+# Template variables:
+#   {cuda_version}   - Full CUDA version (e.g., "12.8")
+#   {cuda_short}     - CUDA without dot (e.g., "128")
+#   {cuda_short2}    - CUDA short for spconv (e.g., "124" not "1240")
+#   {cuda_major}     - CUDA major version (e.g., "12")
+#   {torch_version}  - Full PyTorch version (e.g., "2.8.0")
+#   {torch_short}    - PyTorch without dots (e.g., "280")
+#   {torch_mm}       - PyTorch major.minor no dot (e.g., "28")
+#   {torch_dotted_mm}- PyTorch major.minor with dot (e.g., "2.8")
+#   {py_version}     - Python version (e.g., "3.10")
+#   {py_short}       - Python without dot (e.g., "310")
+#   {py_minor}       - Python minor version only (e.g., "10")
+#   {py_tag}         - Python tag (e.g., "cp310")
+#   {platform}       - Platform tag (e.g., "linux_x86_64")
+#
+# Install methods:
+#   index          - pip --extra-index-url (PEP 503)
+#   find_links     - pip --find-links
+#   github_index   - GitHub Pages index (--find-links)
+#   pypi_variant   - Package name varies by CUDA (e.g., spconv-cu124)
+#   github_release - Direct wheel URL from GitHub releases
+#   pypi           - Standard PyPI install
+
+packages:
+  # ===========================================================================
+  # PyTorch Geometric (PyG) - official index
+  # https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html
+  # ===========================================================================
+  torch-scatter:
+    method: find_links
+    index_url: "https://data.pyg.org/whl/torch-{torch_version}+cu{cuda_short}.html"
+    description: Scatter operations for PyTorch
+
+  torch-cluster:
+    method: find_links
+    index_url: "https://data.pyg.org/whl/torch-{torch_version}+cu{cuda_short}.html"
+    description: Clustering algorithms for PyTorch
+
+  torch-sparse:
+    method: find_links
+    index_url: "https://data.pyg.org/whl/torch-{torch_version}+cu{cuda_short}.html"
+    description: Sparse tensor operations for PyTorch
+
+  torch-spline-conv:
+    method: find_links
+    index_url: "https://data.pyg.org/whl/torch-{torch_version}+cu{cuda_short}.html"
+    description: Spline convolutions for PyTorch
+
+  # ===========================================================================
+  # pytorch3d - Facebook's official wheels
+  # https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md
+  # ===========================================================================
+  pytorch3d:
+    method: index
+    index_url: "https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py3{py_minor}_cu{cuda_short}_pyt{torch_short}/download.html"
+    description: PyTorch3D - 3D deep learning library
+
+  # ===========================================================================
+  # PozzettiAndrea wheel repos (GitHub Pages indexes)
+  # ===========================================================================
+  nvdiffrast:
+    method: github_index
+    index_url: "https://pozzettiandrea.github.io/nvdiffrast-full-wheels/cu{cuda_short}-torch{torch_short}/"
+    description: NVIDIA differentiable rasterizer
+
+  cumesh:
+    method: github_index
+    index_url: "https://pozzettiandrea.github.io/cumesh-wheels/cu{cuda_short}-torch{torch_short}/"
+    description: CUDA-accelerated mesh utilities
+
+  o_voxel:
+    method: github_index
+    index_url: "https://pozzettiandrea.github.io/ovoxel-wheels/cu{cuda_short}-torch{torch_short}/"
+    description: O-Voxel CUDA extension for TRELLIS
+
+  flex_gemm:
+    method: github_index
+    index_url: "https://pozzettiandrea.github.io/flexgemm-wheels/cu{cuda_short}-torch{torch_short}/"
+    description: Flexible GEMM operations
+
+  nvdiffrec_render:
+    method: github_release
+    description: NVDiffRec rendering utilities
+    sources:
+      - name: PozzettiAndrea
+        url_template: "https://github.com/PozzettiAndrea/nvdiffrec_render-wheels/releases/download/cu{cuda_short}-torch{torch_short}/nvdiffrec_render-{version}%2Bcu{cuda_short}torch{torch_mm}-{py_tag}-{py_tag}-linux_x86_64.whl"
+        platforms: [linux_x86_64]
+      - name: PozzettiAndrea-windows
+        url_template: "https://github.com/PozzettiAndrea/nvdiffrec_render-wheels/releases/download/cu{cuda_short}-torch{torch_short}/nvdiffrec_render-{version}%2Bcu{cuda_short}torch{torch_mm}-{py_tag}-{py_tag}-win_amd64.whl"
+        platforms: [win_amd64]
+
+  # ===========================================================================
+  # spconv - PyPI with CUDA-versioned package names
+  # ===========================================================================
+  spconv:
+    method: pypi_variant
+    package_template: "spconv-cu{cuda_short2}"
+    description: Sparse convolution library (use spconv-cu126 for CUDA 12.6+)
+
+  # ===========================================================================
+  # sageattention - Fast quantized attention
+  # ===========================================================================
+  sageattention:
+    method: github_release
+    description: SageAttention - 2-5x faster than FlashAttention with quantized kernels
+    sources:
+      - name: kijai-hf
+        url_template: "https://huggingface.co/Kijai/PrecompiledWheels/resolve/main/sageattention-{version}-cp312-cp312-linux_x86_64.whl"
+        platforms: [linux_x86_64]
+      - name: woct0rdho
+        url_template: "https://github.com/woct0rdho/SageAttention/releases/download/v2.2.0-windows.post3/sageattention-2.2.0%2Bcu{cuda_short}torch{torch_version}.post3-cp39-abi3-win_amd64.whl"
+        platforms: [win_amd64]
+
+  # ===========================================================================
+  # triton - Required for sageattention on Linux
+  # ===========================================================================
+  triton:
+    method: pypi
+    description: Triton compiler for custom CUDA kernels
+
+  # ===========================================================================
+  # flash-attn - Multi-source prebuilt wheels
+  # ===========================================================================
+  flash-attn:
+    method: github_release
+    description: Flash Attention for fast transformer inference
+    sources:
+      - name: Dao-AILab
+        url_template: "https://github.com/Dao-AILab/flash-attention/releases/download/v{version}/flash_attn-{version}%2Bcu{cuda_major}torch{torch_dotted_mm}cxx11abiTRUE-{py_tag}-{py_tag}-linux_x86_64.whl"
+        platforms: [linux_x86_64]
+      - name: mjun0812
+        url_template: "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.7.7/flash_attn-{version}%2Bcu{cuda_short}torch{torch_dotted_mm}-{py_tag}-{py_tag}-linux_x86_64.whl"
+        platforms: [linux_x86_64]
+      - name: mjun0812-windows
+        url_template: "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.7.7/flash_attn-{version}%2Bcu{cuda_short}torch{torch_dotted_mm}-{py_tag}-{py_tag}-win_amd64.whl"
+        platforms: [win_amd64]
+      - name: bdashore3
+        url_template: "https://github.com/bdashore3/flash-attention/releases/download/v{version}/flash_attn-{version}%2Bcu{cuda_short}torch{torch_version}cxx11abiFALSE-{py_tag}-{py_tag}-win_amd64.whl"
+        platforms: [win_amd64]

{comfy_env-0.0.12.dist-info → comfy_env-0.0.13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: comfy-env
-Version: 0.0.12
+Version: 0.0.13
 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
@@ -17,6 +17,7 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Requires-Python: >=3.10
+Requires-Dist: pyyaml>=6.0
 Requires-Dist: tomli>=2.0.0; python_version < '3.11'
 Requires-Dist: uv>=0.4.0
 Provides-Extra: dev

{comfy_env-0.0.12.dist-info → comfy_env-0.0.13.dist-info}/RECORD

@@ -1,15 +1,14 @@
 comfy_env/__init__.py,sha256=1EjSpAK20Z6RAXKn3onZwj2a_FIctG0E_LhWlE7SgIs,4250
-comfy_env/cli.py,sha256=9GvQYrXlJRl0ZaCuFHvRtVxWQ34Axd5Brgu5FWRONp4,11424
+comfy_env/cli.py,sha256=q4y_tlPyqKMZhge7XeO9VdbFVZ4dl9LZsgnnTVQYXD4,15979
 comfy_env/decorator.py,sha256=daFR5aLzshkmo5sRKhSGPcTUgIUWml7Gs6A1bfnDuyc,15718
-comfy_env/errors.py,sha256=egeyXY-j7KpxyA0s67TcJLEfJX23LCAD3v1P4FgQIGE,10917
-comfy_env/install.py,sha256=JaatBM5PyY4Ua6E2HgkLAlFZ2TSEjlI-lAwJ-Hl7wU8,18876
-comfy_env/registry.py,sha256=iZprw2iIAJl-Vvotsf_B_PofzqE-6IiVfNPBHYUYg6g,11577
-comfy_env/resolver.py,sha256=xz7GPlxy02iwwpocIzzbdGnrwnSpi-D5IzpL8SQSgvI,12893
-comfy_env/runner.py,sha256=0YUqzK93u--7pKV6_yVC564AJE9rS3y81t5ZhQi2t4Y,9696
+comfy_env/errors.py,sha256=8hN8NDlo8oBUdapc-eT3ZluigI5VBzfqsSBvQdfWlz4,9943
+comfy_env/install.py,sha256=q8x7sQmCWMQyUjeBcbrg39eAP3ZPars54iepupMe2pE,18922
+comfy_env/registry.py,sha256=uFCtGmWYvwGCqObXgzmArX7o5JsFNsHXxayofk3m6no,2569
+comfy_env/resolver.py,sha256=AC_xizQARVdmc8ZpGrIPQd-oA6Gv8w49OsHKltRAq5o,12251
 comfy_env/env/__init__.py,sha256=sybOBrxJCfL4Xry9NNd5xwn9hXIHudXlXDa7SpJkPCE,811
 comfy_env/env/config.py,sha256=R8JyE5iQLHKgnxXOGA8SAI7iu2eYSfXn-MsaqHoU2_A,5667
-comfy_env/env/config_file.py,sha256=2sTeBUqdDYPwQfUO-3Mu2zTsgxQ4IhllRxEd2a4pgPg,20593
-comfy_env/env/detection.py,sha256=Br_Ed2aqGR6_M01UMZnrzVky3a2iyzz1XuhOXduM1zw,5993
+comfy_env/env/config_file.py,sha256=6Mh2VsNvyFJuTL4_vOX0sG3Q5SMHLI1P97aY8FFvwnk,20696
+comfy_env/env/detection.py,sha256=L4kXKGs9j7Vw4Qeh6Zw5PfVZW0dN2zuJqojDPgINIhA,5653
 comfy_env/env/manager.py,sha256=MRmqg3Qw0IYUwR6FtM3ZIcGZvWcnKDyO8xbzIAxnwCk,24782
 comfy_env/env/security.py,sha256=dNSitAnfBNVdvxgBBntYw33AJaCs_S1MHb7KJhAVYzM,8171
 comfy_env/env/platform/__init__.py,sha256=Nb5MPZIEeanSMEWwqU4p4bnEKTJn1tWcwobnhq9x9IY,614
@@ -18,7 +17,7 @@ comfy_env/env/platform/darwin.py,sha256=HK3VkLT6DfesAnIXwx2IaUFHTBclF0xTQnC7azWY
 comfy_env/env/platform/linux.py,sha256=xLp8FEbFqZLQrzIZBI9z3C4g23Ab1ASTHLsXDzsdCoA,2062
 comfy_env/env/platform/windows.py,sha256=nD1-bKU2rGmEJlS-cc5yWXMSA51YQtVupn-lQEO5UYA,14840
 comfy_env/ipc/__init__.py,sha256=pTjgJn5YJxLXmEvuKh3lkCEJQs-6W6_F01jfkFMUi0c,1375
-comfy_env/ipc/bridge.py,sha256=kEy__kco8FVQNj5MyadF5k00YEivcGmifAJAOfr643U,17645
+comfy_env/ipc/bridge.py,sha256=zcyN3xzV4WWBrBFNwCniPBR58dLCg46-k9TtyW5U000,16437
 comfy_env/ipc/protocol.py,sha256=gfWe5yEDUn4QWhcdWFcxn40GqxlW1Uf23j0edOzPPng,7951
 comfy_env/ipc/tensor.py,sha256=DyU28GymKkLPVwzZyKdm2Av222hdaycMgv3KdL5mtO0,12009
 comfy_env/ipc/torch_bridge.py,sha256=WzdwDJa3N_1fEl9OeZxikvMbwryO5P63o0WmEDpS18A,13206
@@ -32,8 +31,9 @@ comfy_env/workers/pool.py,sha256=MtjeOWfvHSCockq8j1gfnxIl-t01GSB79T5N4YB82Lg,695
 comfy_env/workers/tensor_utils.py,sha256=TCuOAjJymrSbkgfyvcKtQ_KbVWTqSwP9VH_bCaFLLq8,6409
 comfy_env/workers/torch_mp.py,sha256=DsfxE3LBAWEuGtk-p-YL0UhVQ7VDh73KT_TFRxYN4-Q,12563
 comfy_env/workers/venv.py,sha256=_ekHfZPqBIPY08DjqiXm6cTBQH4DrbxRWR3AAv3mit8,31589
-comfy_env-0.0.12.dist-info/METADATA,sha256=h4b6TPaQ2ubWFwkDRynrbn9Ccx3Nz2vEr6TaguVbRQE,5372
-comfy_env-0.0.12.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-comfy_env-0.0.12.dist-info/entry_points.txt,sha256=J4fXeqgxU_YenuW_Zxn_pEL7J-3R0--b6MS5t0QmAr0,49
-comfy_env-0.0.12.dist-info/licenses/LICENSE,sha256=E68QZMMpW4P2YKstTZ3QU54HRQO8ecew09XZ4_Vn870,1093
-comfy_env-0.0.12.dist-info/RECORD,,
+comfy_env/wheel_sources.yml,sha256=ubVuQllCQGkZhLNQaG54divCwn0zLzYg4turzhnIZQ8,7150
+comfy_env-0.0.13.dist-info/METADATA,sha256=FRLYYBNxWwLm5I3vWRVxT1LYp5wvW13ibZ95PVdOaU4,5399
+comfy_env-0.0.13.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+comfy_env-0.0.13.dist-info/entry_points.txt,sha256=J4fXeqgxU_YenuW_Zxn_pEL7J-3R0--b6MS5t0QmAr0,49
+comfy_env-0.0.13.dist-info/licenses/LICENSE,sha256=E68QZMMpW4P2YKstTZ3QU54HRQO8ecew09XZ4_Vn870,1093
+comfy_env-0.0.13.dist-info/RECORD,,

comfy_env/runner.py

@@ -1,273 +0,0 @@
-"""
-Generic runner for isolated subprocess execution.
-
-This module is the entry point for subprocess execution. The runner handles
-requests for ANY @isolated class in the environment, importing classes on demand.
-
-Usage (Unix Domain Socket - recommended):
-    python -m comfy_env.runner \
-        --node-dir /path/to/ComfyUI-SAM3DObjects/nodes \
-        --comfyui-base /path/to/ComfyUI \
-        --import-paths ".,../vendor" \
-        --socket /tmp/comfyui-isolation-myenv-12345.sock
-
-Usage (Legacy stdin/stdout):
-    python -m comfy_env.runner \
-        --node-dir /path/to/ComfyUI-SAM3DObjects/nodes \
-        --comfyui-base /path/to/ComfyUI \
-        --import-paths ".,../vendor"
-
-The runner:
-1. Sets COMFYUI_ISOLATION_WORKER=1 (so @isolated decorator becomes no-op)
-2. Adds paths to sys.path
-3. Connects to Unix Domain Socket (or uses stdin/stdout)
-4. Dynamically imports classes as needed (cached)
-5. Calls methods and returns responses
-"""
-
-import os
-import sys
-import json
-import argparse
-import traceback
-import warnings
-import logging
-import importlib
-from typing import Any, Dict, Optional
-
-# Suppress warnings that could interfere with JSON IPC
-warnings.filterwarnings("ignore")
-os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")
-logging.disable(logging.WARNING)
-
-# Mark that we're in worker mode - this makes @isolated decorator a no-op
-os.environ["COMFYUI_ISOLATION_WORKER"] = "1"
-
-
-def setup_paths(node_dir: str, comfyui_base: Optional[str], import_paths: Optional[str]):
-    """Setup sys.path for imports."""
-    from pathlib import Path
-
-    node_path = Path(node_dir)
-
-    # Set COMFYUI_BASE env var for stubs to use
-    if comfyui_base:
-        os.environ["COMFYUI_BASE"] = comfyui_base
-
-    # Add comfyui-isolation stubs directory (provides folder_paths, etc.)
-    stubs_dir = Path(__file__).parent / "stubs"
-    sys.path.insert(0, str(stubs_dir))
-
-    # Add import paths
-    if import_paths:
-        for p in import_paths.split(","):
-            p = p.strip()
-            if p:
-                full_path = node_path / p
-                sys.path.insert(0, str(full_path))
-
-    # Add node_dir itself
-    sys.path.insert(0, str(node_path))
-
-
-def serialize_result(obj: Any) -> Any:
-    """Serialize result for JSON transport."""
-    from comfy_env.ipc.protocol import encode_object
-    return encode_object(obj)
-
-
-def deserialize_arg(obj: Any) -> Any:
-    """Deserialize argument from JSON transport."""
-    from comfy_env.ipc.protocol import decode_object
-    return decode_object(obj)
-
-
-# Cache for imported classes and instances
-_class_cache: Dict[str, type] = {}
-_instance_cache: Dict[str, object] = {}
-
-
-def get_instance(module_name: str, class_name: str) -> object:
-    """Get or create an instance of a class."""
-    cache_key = f"{module_name}.{class_name}"
-
-    if cache_key not in _instance_cache:
-        # Import the class if not cached
-        if cache_key not in _class_cache:
-            print(f"[Runner] Importing {class_name} from {module_name}...", file=sys.stderr)
-            module = importlib.import_module(module_name)
-            cls = getattr(module, class_name)
-            _class_cache[cache_key] = cls
-
-        # Create instance
-        cls = _class_cache[cache_key]
-        _instance_cache[cache_key] = cls()
-        print(f"[Runner] Created instance of {class_name}", file=sys.stderr)
-
-    return _instance_cache[cache_key]
-
-
-def run_worker(
-    node_dir: str,
-    comfyui_base: Optional[str],
-    import_paths: Optional[str],
-    socket_path: Optional[str] = None,
-):
-    """
-    Main worker loop - handles JSON-RPC requests via transport.
-
-    Args:
-        node_dir: Path to node package directory
-        comfyui_base: Path to ComfyUI base directory
-        import_paths: Comma-separated import paths
-        socket_path: Unix domain socket path (if None, uses stdin/stdout)
-    """
-    from comfy_env.ipc.transport import UnixSocketTransport, StdioTransport
-
-    # Setup paths first
-    setup_paths(node_dir, comfyui_base, import_paths)
-
-    # Create transport
-    if socket_path:
-        # Unix Domain Socket transport (recommended)
-        print(f"[Runner] Connecting to socket: {socket_path}", file=sys.stderr)
-        transport = UnixSocketTransport.connect(socket_path)
-        use_uds = True
-    else:
-        # Legacy stdin/stdout transport
-        print("[Runner] Using stdin/stdout transport", file=sys.stderr)
-        transport = StdioTransport()
-        use_uds = False
-
-    try:
-        # Send ready signal
-        transport.send({"status": "ready"})
-
-        # Main loop - read requests, execute, respond
-        while True:
-            response = {"jsonrpc": "2.0", "id": None}
-
-            try:
-                request = transport.recv()
-                response["id"] = request.get("id")
-
-                method_name = request.get("method")
-                params = request.get("params", {})
-
-                if method_name == "shutdown":
-                    # Clean shutdown
-                    response["result"] = {"status": "shutdown"}
-                    transport.send(response)
-                    break
-
-                # Get module/class from request
-                module_name = request.get("module")
-                class_name = request.get("class")
-
-                if not module_name or not class_name:
-                    response["error"] = {
-                        "code": -32602,
-                        "message": "Missing 'module' or 'class' in request",
-                    }
-                    transport.send(response)
-                    continue
-
-                # Get or create instance
-                try:
-                    instance = get_instance(module_name, class_name)
-                except Exception as e:
-                    response["error"] = {
-                        "code": -32000,
-                        "message": f"Failed to import {module_name}.{class_name}: {e}",
-                        "data": {"traceback": traceback.format_exc()}
-                    }
-                    transport.send(response)
-                    continue
-
-                # Get the method
-                method = getattr(instance, method_name, None)
-                if method is None:
-                    response["error"] = {
-                        "code": -32601,
-                        "message": f"Method not found: {method_name}",
-                    }
-                    transport.send(response)
-                    continue
-
-                # Deserialize arguments
-                deserialized_params = {}
-                for key, value in params.items():
-                    deserialized_params[key] = deserialize_arg(value)
-
-                # For legacy stdio transport, redirect stdout to stderr during execution
-                # This prevents print() in node code from corrupting JSON protocol
-                # (UDS transport doesn't need this since it uses a separate socket)
-                if not use_uds:
-                    original_stdout = sys.stdout
-                    sys.stdout = sys.stderr
-
-                    # Also redirect at file descriptor level for C libraries
-                    stdout_fd = original_stdout.fileno()
-                    stderr_fd = sys.stderr.fileno()
-                    stdout_fd_copy = os.dup(stdout_fd)
-                    os.dup2(stderr_fd, stdout_fd)
-
-                # Call the method
-                print(f"[Runner] Calling {class_name}.{method_name}...", file=sys.stderr)
-                try:
-                    result = method(**deserialized_params)
-                finally:
-                    if not use_uds:
-                        # Restore file descriptor first, then Python stdout
-                        os.dup2(stdout_fd_copy, stdout_fd)
-                        os.close(stdout_fd_copy)
-                        sys.stdout = original_stdout
-
-                # Serialize result
-                serialized_result = serialize_result(result)
-                response["result"] = serialized_result
-
-                print(f"[Runner] {class_name}.{method_name} completed", file=sys.stderr)
-
-            except ConnectionError as e:
-                # Socket closed - normal shutdown
-                print(f"[Runner] Connection closed: {e}", file=sys.stderr)
-                break
-            except Exception as e:
-                tb = traceback.format_exc()
-                print(f"[Runner] Error: {e}", file=sys.stderr)
-                print(tb, file=sys.stderr)
-                response["error"] = {
-                    "code": -32000,
-                    "message": str(e),
-                    "data": {"traceback": tb}
-                }
-
-            try:
-                transport.send(response)
-            except ConnectionError:
-                break
-
-    finally:
-        transport.close()
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Isolated node runner")
-    parser.add_argument("--node-dir", required=True, help="Node package directory")
-    parser.add_argument("--comfyui-base", help="ComfyUI base directory")
-    parser.add_argument("--import-paths", help="Comma-separated import paths")
-    parser.add_argument("--socket", help="Unix domain socket path (if not provided, uses stdin/stdout)")
-
-    args = parser.parse_args()
-
-    run_worker(
-        node_dir=args.node_dir,
-        comfyui_base=args.comfyui_base,
-        import_paths=args.import_paths,
-        socket_path=args.socket,
-    )
-
-
-if __name__ == "__main__":
-    main()