custom comfyui venv (#117)

* it works

* print

* refact

* refact

* more logging

* no f

* Update comfyui_process.py

install_comfyui.py

@@ -1,12 +1,12 @@
 import os
 import sys
-import git
 
 
 default_install_location = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'ComfyUI')
 
 
 def main(install_location):
+    import git
     git_repo_url = 'https://github.com/comfyanonymous/ComfyUI.git'
     os.mkdir(install_location)
     git.Repo.clone_from(git_repo_url, install_location)

lib_comfyui/comfyui/pre_main.py

@@ -20,11 +20,35 @@ 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...')

lib_comfyui/comfyui/print_sys_path.py

@@ -0,0 +1,6 @@
+import sys
+import os
+
+
+if __name__ == '__main__':
+    print(os.pathsep.join(sys.path), end='')

lib_comfyui/comfyui_process.py

@@ -3,9 +3,11 @@
 import os
 import subprocess
 import sys
+from pathlib import Path
+
 from lib_comfyui import ipc, torch_utils, argv_conversion, global_state
 from lib_comfyui.webui import settings
-from lib_comfyui.comfyui import pre_main
+from lib_comfyui.comfyui import pre_main, print_sys_path
 
 
 comfyui_process = None
@@ -19,7 +21,7 @@ def start():
         return
 
     install_location = settings.get_install_location()
-    if not os.path.exists(install_location):
+    if not install_location.exists():
         print('[sd-webui-comfyui]', f'Could not find ComfyUI under directory "{install_location}". The server will NOT be started.', file=sys.stderr)
         return
 
@@ -34,28 +36,74 @@ def start():
 def start_comfyui_process(comfyui_install_location):
     global comfyui_process
 
-    comfyui_env = os.environ.copy()
-    comfyui_sys_path = get_base_sys_path()
-    comfyui_sys_path.insert(1, settings.get_extension_base_dir())
-    comfyui_env['PYTHONPATH'] = os.pathsep.join(comfyui_sys_path)
-    comfyui_env['SD_WEBUI_COMFYUI_IPC_STRATEGY_CLASS_NAME'] = global_state.ipc_strategy_class.__name__
-
-    args = [sys.executable, inspect.getfile(pre_main)] + argv_conversion.get_comfyui_args()
-
+    executable = get_comfyui_executable(comfyui_install_location)
+    comfyui_env = get_comfyui_env(executable, 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(
         args=args,
-        executable=sys.executable,
-        cwd=comfyui_install_location,
+        executable=executable,
+        cwd=str(comfyui_install_location),
         env=comfyui_env,
     )
 
 
-def get_base_sys_path():
+def get_comfyui_executable(comfyui_install_location):
+    executable = sys.executable
+    venv = comfyui_install_location / 'venv'
+    if venv.exists():
+        if os.name == 'nt':
+            executable = venv / 'scripts' / 'python.exe'
+        else:
+            executable = venv / 'bin' / 'python'
+
+        print('[sd-webui-comfyui]', 'Detected custom ComfyUI venv:', venv)
+
+    return str(executable)
+
+
+def get_comfyui_env(executable, 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)
+    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()
+    del env['PYTHONPATH']
     return subprocess.run(
-        [sys.executable, '-c', "import sys, os; print(os.pathsep.join(sys.path))"],
+        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)
+    ).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
+        return
+
+    print('[sd-webui-comfyui]', 'Installing mandatory pip requirements in ComfyUI venv...')
+    subprocess.check_call(
+        args=[
+            executable,
+            *(['-s'] if "python_embeded" in executable or "python_embedded" in executable else []),
+            '-m',
+            'pip',
+            'install',
+            '-r',
+            str(Path(settings.get_extension_base_dir(), 'requirements.txt')),
+        ],
+        executable=executable,
+        cwd=str(comfyui_install_location),
+        env=comfyui_env,
+    )
 
 
 @ipc.restrict_to_process('webui')

lib_comfyui/webui/settings.py

@@ -1,5 +1,6 @@
 import sys
 import textwrap
+from pathlib import Path
 from lib_comfyui import ipc, global_state
 import install_comfyui
 
@@ -68,11 +69,11 @@ def update_comfyui_graceful_termination_timeout():
 
 
 @ipc.restrict_to_process('webui')
-def get_install_location():
+def get_install_location() -> Path:
     from modules import shared
     install_location = install_comfyui.default_install_location
     install_location = shared.opts.data.get('comfyui_install_location', install_location).strip()
-    return install_location
+    return Path(install_location)
 
 
 @ipc.restrict_to_process('webui')