Fixes issues with running from a vanilla ComfyUI workspace and having this package installed (fixes #30)

comfy/__init__.py

@@ -1,3 +1,8 @@
 # This file is automatically generated by the build process when version is
 # updated in pyproject.toml.
 __version__ = "0.3.56"
+
+# This deals with workspace issues
+from comfy_compatibility.workspace import auto_patch_workspace_and_restart
+
+auto_patch_workspace_and_restart()

comfy_compatibility/imports.py

@@ -0,0 +1,77 @@
+import os
+import sys
+import inspect
+from importlib.machinery import PathFinder
+
+MAIN_PY = 'main_py'
+CURRENT_DIRECTORY = 'cwd'
+SITE_PACKAGES = 'site'
+
+
+class ImportContext:
+    def __init__(self, *module_names, order):
+        self.module_names = module_names
+        self.order = order
+        self.original_modules = {}
+        self.finder = None
+
+    def __enter__(self):
+        try:
+            main_frame = next(f for f in reversed(inspect.stack()) if f.frame.f_globals.get('__name__') == '__main__')
+            self.main_py_dir = os.path.dirname(os.path.abspath(main_frame.filename))
+        except (StopIteration, AttributeError):
+            self.main_py_dir = None
+
+        self.original_modules = {name: sys.modules.get(name) for name in self.module_names}
+        for name in self.module_names:
+            if name in sys.modules:
+                for mod_name in list(sys.modules.keys()):
+                    if mod_name == name or mod_name.startswith(f"{name}."):
+                        del sys.modules[mod_name]
+
+        self.finder = self._CustomFinder(self.module_names, self.order, self.main_py_dir)
+        sys.meta_path.insert(0, self.finder)
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.finder in sys.meta_path:
+            sys.meta_path.remove(self.finder)
+
+        for name in self.module_names:
+            for mod_name in list(sys.modules.keys()):
+                if mod_name == name or mod_name.startswith(f"{name}."):
+                    del sys.modules[mod_name]
+
+        for name, mod in self.original_modules.items():
+            if mod:
+                sys.modules[name] = mod
+            elif name in sys.modules:
+                del sys.modules[name]
+
+    class _CustomFinder(PathFinder):
+        def __init__(self, names, order, main_py_dir):
+            self.module_names = set(names)
+            self.order = order
+            self.main_py_dir = main_py_dir
+
+        def find_spec(self, fullname, path=None, target=None):
+            if fullname not in self.module_names:
+                return None
+
+            cwd_path = os.getcwd()
+            site_paths = [p for p in sys.path if 'site-packages' in p]
+            other_paths = [p for p in sys.path if p != cwd_path and p not in site_paths]
+
+            search_path = []
+            for source in self.order:
+                if source == CURRENT_DIRECTORY:
+                    search_path.append(cwd_path)
+                elif source == MAIN_PY and self.main_py_dir:
+                    search_path.append(self.main_py_dir)
+                elif source == SITE_PACKAGES:
+                    search_path.extend(site_paths)
+
+            search_path.extend(other_paths)
+
+            return super().find_spec(fullname, path=search_path, target=target)

comfy_compatibility/workspace.py

@@ -0,0 +1,84 @@
+import inspect
+import logging
+import os
+import shutil
+import sys
+
+# mitigations for comfyui workspace chaos
+logger = logging.getLogger(__name__)
+
+
+def auto_patch_workspace_and_restart():
+    """
+    Detects a specific workspace structure, creates necessary __init__.py files
+    to make directories proper Python packages, and then restarts the application.
+    If the workspace is a Git repository, it locally ignores the created files.
+    """
+    try:
+        main_frame = next(f for f in reversed(inspect.stack()) if f.frame.f_globals.get('__name__') == '__main__')
+        main_file_path = main_frame.filename
+        workspace_dir = os.path.dirname(os.path.abspath(main_file_path))
+    except (StopIteration, AttributeError):
+        return
+
+    if not os.path.isfile(os.path.join(workspace_dir, 'nodes.py')):
+        return
+
+    git_is_available = False
+    if shutil.which('git') and os.path.isdir(os.path.join(workspace_dir, '.git')):
+        git_is_available = True
+
+    target_base_dirs = [
+        'comfy',
+        'comfy_extras',
+        'comfy_execution',
+        'comfy_api',
+        'comfy_config'
+    ]
+
+    patched_any_file = False
+
+    for dir_name in target_base_dirs:
+        start_dir = os.path.join(workspace_dir, dir_name)
+        if not os.path.isdir(start_dir):
+            continue
+
+        for dirpath, _, filenames in os.walk(start_dir):
+            init_py_path = os.path.join(dirpath, '__init__.py')
+
+            if os.path.exists(init_py_path):
+                continue
+
+            if any(fname.endswith('.py') for fname in filenames):
+                logger.debug(f"  Initializing package: {dirpath}")
+                try:
+                    with open(init_py_path, 'w') as f:
+                        pass
+                    patched_any_file = True
+
+                    if git_is_available:
+                        try:
+                            relative_path = os.path.relpath(init_py_path, workspace_dir).replace(os.sep, '/')
+
+                            exclude_file = os.path.join(workspace_dir, '.git', 'info', 'exclude')
+                            os.makedirs(os.path.dirname(exclude_file), exist_ok=True)
+
+                            content = ""
+                            if os.path.exists(exclude_file):
+                                with open(exclude_file, 'r') as f_read:
+                                    content = f_read.read()
+
+                            if relative_path not in content.splitlines():
+                                with open(exclude_file, 'a') as f_append:
+                                    f_append.write(f"\n{relative_path}")
+                                logger.debug(f"  Ignoring via .git/info/exclude: {relative_path}")
+
+                        except Exception as e:
+                            logger.debug(f"Warning: Could not add {relative_path} to Git exclude file. Error: {e}")
+
+                except OSError as e:
+                    logger.debug(f"Warning: Could not create {init_py_path}. Error: {e}")
+
+    if patched_any_file:
+        logger.debug("Found and initialized Python package directories in your workspace. This is a one-time operation to enable proper imports. Now restarting...")
+        os.execv(sys.executable, [sys.executable] + sys.argv)

pyproject.toml

@@ -277,4 +277,4 @@ exclude = ["*.ipynb", "**/generated/*.pyi"]
 allow-direct-references = true
 
 [tool.hatch.build.targets.wheel]
-packages = ["comfy/", "comfy_extras/", "comfy_api/", "comfy_api_nodes/", "comfy_config/", "comfy_execution/"]
+packages = ["comfy/", "comfy_extras/", "comfy_api/", "comfy_api_nodes/", "comfy_config/", "comfy_execution/", "comfy_compatibility/"]

tests/issues/__test_30_comfy_workspace_conflicts.py

@@ -0,0 +1,18 @@
+"""
+# get the workspace and set it up
+git clone git@github.com:comfyanonymous/ComfyUI.git comfyanonymous_ComfyUI
+cd comfyanonymous_ComfyUI
+uv venv
+source .venv/bin/activate
+uv pip install --torch-backend=auto requirements.txt
+
+# install livepeer with up to date comfyui lts
+uv pip install --torch-backend=auto git+https://github.com/livepeer/comfystream.git --overrides=<(echo "comfyui@git+https://github.com/hiddenswitch/ComfyUI.git@fixes/issue-30")
+
+# install the nodes
+curl -L -o comfystream.zip https://github.com/livepeer/comfystream/archive/refs/heads/main.zip
+mkdir -p custom_nodes
+unzip comfystream.zip -d custom_nodes/
+rm comfystream.zip
+
+"""