From 9b04919330486e301760f7adcacc6df30e5a31cf Mon Sep 17 00:00:00 2001
From: doctorpangloss <@hiddenswitch.com>
Date: Thu, 4 Sep 2025 14:02:10 -0700
Subject: [PATCH] When comfy_execution is touched, this is also a good chance
 to inject __init__.py files

---
 comfy_compatibility/workspace.py | 31 +++++++++++++++++--------------
 comfy_execution/__init__.py      |  4 ++++
 2 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/comfy_compatibility/workspace.py b/comfy_compatibility/workspace.py
index 8bd8d95f8..26070cdbb 100644
--- a/comfy_compatibility/workspace.py
+++ b/comfy_compatibility/workspace.py
@@ -3,25 +3,29 @@ 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():
+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.
     """
-    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 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
@@ -65,14 +69,14 @@ def auto_patch_workspace_and_restart():
         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}")
+                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"Warning: Could not create {init_py_path}. Error: {e}")
+                    logger.debug(f"could not create {init_py_path}. Error: {e}")
 
     if git_is_available and files_to_ignore:
         try:
@@ -90,12 +94,11 @@ def auto_patch_workspace_and_restart():
                     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}")
+                        logger.debug(f"ignoring via .git/info/exclude: {relative_path}")
 
         except Exception as e:
-            logger.debug(f"Warning: Could not update Git exclude file. Error: {e}")
+            logger.debug(f"could not update Git exclude file. 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...")
+    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)
-
diff --git a/comfy_execution/__init__.py b/comfy_execution/__init__.py
index e69de29bb..390bba771 100644
--- a/comfy_execution/__init__.py
+++ b/comfy_execution/__init__.py
@@ -0,0 +1,4 @@
+# There are some circumstances where comfy/__init__.py may already exist, so let's patch it here too
+from comfy_compatibility.workspace import auto_patch_workspace_and_restart
+
+auto_patch_workspace_and_restart()