From cf7b9f2be2fd2373238afac7b0f53b60493078f8 Mon Sep 17 00:00:00 2001
From: ljleb <lebel.louisjacob@gmail.com>
Date: Thu, 31 Aug 2023 02:16:12 -0400
Subject: [PATCH] patch sys.path inside the comfyui process

---
 lib_comfyui/comfyui/pre_main.py | 40 ++++++++++++---------------------
 lib_comfyui/comfyui_process.py  | 26 ++++++---------------
 2 files changed, 21 insertions(+), 45 deletions(-)

diff --git a/lib_comfyui/comfyui/pre_main.py b/lib_comfyui/comfyui/pre_main.py
index e1ef09d..d4b7526 100644
--- a/lib_comfyui/comfyui/pre_main.py
+++ b/lib_comfyui/comfyui/pre_main.py
@@ -1,8 +1,20 @@
+import os
+import sys
+
+if __name__ == "__main__":
+    # patch sys.path
+    install_dir = os.getenv("SD_WEBUI_COMFYUI_INSTALL_DIR")
+    extension_dir = os.getenv("SD_WEBUI_COMFYUI_EXTENSION_DIR")
+    if not install_dir or not extension_dir:
+        print("[sd-webui-comfyui]", f"Could not add new entries to sys.path. install_dir={install_dir}, extension_dir={extension_dir}, sys.path={sys.path}", file=sys.stderr)
+        print("[sd-webui-comfyui]", f"Exiting...", file=sys.stderr)
+        exit(1)
+
+    sys.path[:0] = (install_dir, extension_dir)
+
 import atexit
 import builtins
 import signal
-import sys
-import os
 import runpy
 from lib_comfyui import (
     custom_extension_injector,
@@ -20,35 +32,11 @@ def comfyui_print(*args, **kwargs):
 @ipc.restrict_to_process('comfyui')
 def main():
     builtins.print = comfyui_print
-    fix_path()
     setup_ipc()
     patch_comfyui()
     start_comfyui()
 
 
-@ipc.restrict_to_process('comfyui')
-def fix_path():
-    def make_path_unique():
-        path = sys.path.copy()
-        sys.path.clear()
-        seen = set()
-        sys.path.extend(
-            p for p in path
-            if not (p in seen or seen.add(p))
-        )
-
-    def move_comfyui_to_front():
-        comfyui_dir = os.getcwd()
-        try:
-            sys.path.remove(comfyui_dir)
-        except ValueError:
-            pass
-        sys.path.insert(0, comfyui_dir)
-
-    make_path_unique()
-    move_comfyui_to_front()
-
-
 @ipc.restrict_to_process('comfyui')
 def setup_ipc():
     print('[sd-webui-comfyui]', 'Setting up IPC...')
diff --git a/lib_comfyui/comfyui_process.py b/lib_comfyui/comfyui_process.py
index b6bee89..ff092c4 100644
--- a/lib_comfyui/comfyui_process.py
+++ b/lib_comfyui/comfyui_process.py
@@ -35,7 +35,7 @@ def start_comfyui_process(comfyui_install_location):
     global comfyui_process
 
     executable = get_comfyui_executable(comfyui_install_location)
-    comfyui_env = get_comfyui_env(executable, comfyui_install_location)
+    comfyui_env = get_comfyui_env(comfyui_install_location)
     install_comfyui_requirements(executable, comfyui_install_location, comfyui_env)
     args = [executable, inspect.getfile(pre_main)] + argv_conversion.get_comfyui_args()
     comfyui_process = subprocess.Popen(
@@ -60,29 +60,17 @@ def get_comfyui_executable(comfyui_install_location):
     return str(executable)
 
 
-def get_comfyui_env(executable, comfyui_install_location):
+def get_comfyui_env(comfyui_install_location):
     comfyui_env = os.environ.copy()
-    comfyui_sys_path = get_base_sys_path(executable, comfyui_install_location)
-    comfyui_sys_path[:0] = (str(comfyui_install_location), settings.get_extension_base_dir())
-    comfyui_env['PYTHONPATH'] = os.pathsep.join(comfyui_sys_path)
+    if 'PYTHONPATH' in comfyui_env:
+        del comfyui_env['PYTHONPATH']
+
+    comfyui_env['SD_WEBUI_COMFYUI_INSTALL_DIR'] = str(comfyui_install_location)
+    comfyui_env['SD_WEBUI_COMFYUI_EXTENSION_DIR'] = settings.get_extension_base_dir()
     comfyui_env['SD_WEBUI_COMFYUI_IPC_STRATEGY_CLASS_NAME'] = global_state.ipc_strategy_class.__name__
     return comfyui_env
 
 
-def get_base_sys_path(executable, comfyui_install_location):
-    env = os.environ.copy()
-    if 'PYTHONPATH' in env:
-        del env['PYTHONPATH']
-    return subprocess.run(
-        args=[executable, inspect.getfile(print_sys_path)],
-        executable=executable,
-        cwd=str(comfyui_install_location),
-        env=env,
-        text=True,
-        capture_output=True,
-    ).stdout.split(os.pathsep)[1:]  # remove PYTHONHOME because it is automatically added
-
-
 def install_comfyui_requirements(executable, comfyui_install_location, comfyui_env):
     if executable == sys.executable:
         # requirements already installed in the webui by install.py