mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-09 13:50:49 +08:00
105 lines
4.1 KiB
Python
105 lines
4.1 KiB
Python
import inspect
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import sys
|
|
from typing import Optional
|
|
|
|
# mitigations for comfyui workspace chaos
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def auto_patch_workspace_and_restart(workspace_dir: Optional[str] = None, skip_restart=False):
|
|
"""
|
|
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.
|
|
"""
|
|
if workspace_dir is None:
|
|
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):
|
|
logger.debug("workspace_dir was none and comfyui's main.py is not being run, skipping injection of __init__.py files")
|
|
return
|
|
|
|
if not os.path.isfile(os.path.join(workspace_dir, 'nodes.py')):
|
|
logger.debug(f"did not find a nodes.py inside the workspace_dir {workspace_dir}, skipping injection of __init__.py files")
|
|
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
|
|
files_to_ignore = []
|
|
|
|
for dir_name in target_base_dirs:
|
|
start_dir = os.path.join(workspace_dir, dir_name)
|
|
if not os.path.isdir(start_dir):
|
|
continue
|
|
|
|
dirs_with_py_files = set()
|
|
for dirpath, _, filenames in os.walk(start_dir):
|
|
if any(fname.endswith('.py') for fname in filenames):
|
|
dirs_with_py_files.add(dirpath)
|
|
|
|
if not dirs_with_py_files:
|
|
continue
|
|
|
|
dirs_to_initialize = set()
|
|
for dirpath in dirs_with_py_files:
|
|
parent = dirpath
|
|
while len(parent) >= len(start_dir):
|
|
dirs_to_initialize.add(parent)
|
|
new_parent = os.path.dirname(parent)
|
|
if new_parent == parent:
|
|
break
|
|
parent = new_parent
|
|
|
|
for dirpath in sorted(list(dirs_to_initialize)):
|
|
init_py_path = os.path.join(dirpath, '__init__.py')
|
|
if not os.path.exists(init_py_path):
|
|
logger.debug(f"initializing package: {dirpath}")
|
|
try:
|
|
with open(init_py_path, 'w') as f:
|
|
pass
|
|
patched_any_file = True
|
|
files_to_ignore.append(init_py_path)
|
|
except OSError as e:
|
|
logger.debug(f"could not create {init_py_path}. Error: {e}")
|
|
|
|
if git_is_available and files_to_ignore:
|
|
try:
|
|
exclude_file = os.path.join(workspace_dir, '.git', 'info', 'exclude')
|
|
os.makedirs(os.path.dirname(exclude_file), exist_ok=True)
|
|
|
|
existing_lines = set()
|
|
if os.path.exists(exclude_file):
|
|
with open(exclude_file, 'r') as f_read:
|
|
existing_lines = set(line.strip() for line in f_read)
|
|
|
|
with open(exclude_file, 'a') as f_append:
|
|
for init_py_path in files_to_ignore:
|
|
relative_path = os.path.relpath(init_py_path, workspace_dir).replace(os.sep, '/')
|
|
if relative_path not in existing_lines:
|
|
f_append.write(f"\n{relative_path}")
|
|
existing_lines.add(relative_path)
|
|
logger.debug(f"ignoring via .git/info/exclude: {relative_path}")
|
|
|
|
except Exception as e:
|
|
logger.debug(f"could not update Git exclude file. Error: {e}")
|
|
|
|
if not skip_restart and patched_any_file:
|
|
logger.info("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)
|