From 6f4e2017f4f4f665b559242a138dcf8cb21a40bb Mon Sep 17 00:00:00 2001 From: cmdr2 Date: Mon, 7 Oct 2024 11:28:28 +0530 Subject: [PATCH] Install Forge automatically by creating a conda environment and cloing the forge repo --- ui/easydiffusion/app.py | 33 +++- ui/easydiffusion/backends/ed_classic.py | 1 + ui/easydiffusion/backends/ed_diffusers.py | 1 + ui/easydiffusion/backends/sdkit_common.py | 4 + ui/easydiffusion/backends/webui/__init__.py | 194 ++++++++++++++++++-- ui/easydiffusion/task_manager.py | 3 +- 6 files changed, 207 insertions(+), 29 deletions(-) diff --git a/ui/easydiffusion/app.py b/ui/easydiffusion/app.py index c570c773..6b1d2a09 100644 --- a/ui/easydiffusion/app.py +++ b/ui/easydiffusion/app.py @@ -331,16 +331,31 @@ def open_browser(): webbrowser.open(f"http://localhost:{port}") - Console().print( - Panel( - "\n" - + "[white]Easy Diffusion is ready to serve requests.\n\n" - + "A new browser tab should have been opened by now.\n" - + f"If not, please open your web browser and navigate to [bold yellow underline]http://localhost:{port}/\n", - title="Easy Diffusion is ready", - style="bold yellow on blue", + from easydiffusion.backend_manager import backend + + if backend.is_installed(): + Console().print( + Panel( + "\n" + + "[white]Easy Diffusion is ready to serve requests.\n\n" + + "A new browser tab should have been opened by now.\n" + + f"If not, please open your web browser and navigate to [bold yellow underline]http://localhost:{port}/\n", + title="Easy Diffusion is ready", + style="bold yellow on blue", + ) + ) + else: + backend_name = config["backend"] + Console().print( + Panel( + "\n" + + f"[white]Backend: {backend_name} is still installing..\n\n" + + "Please wait until it finishes installing before making images.\n" + + f"The UI will turn green in the top-right corner once it is ready (in [bold yellow underline]http://localhost:{port}[/]).\n", + title=f"Backend engine is installing", + style="bold yellow on blue", + ) ) - ) def fail_and_die(fail_type: str, data: str): diff --git a/ui/easydiffusion/backends/ed_classic.py b/ui/easydiffusion/backends/ed_classic.py index c9cf745e..07077b6c 100644 --- a/ui/easydiffusion/backends/ed_classic.py +++ b/ui/easydiffusion/backends/ed_classic.py @@ -3,6 +3,7 @@ from sdkit_common import ( stop_backend, install_backend, uninstall_backend, + is_installed, create_sdkit_context, ping, load_model, diff --git a/ui/easydiffusion/backends/ed_diffusers.py b/ui/easydiffusion/backends/ed_diffusers.py index c905652d..b5cf8744 100644 --- a/ui/easydiffusion/backends/ed_diffusers.py +++ b/ui/easydiffusion/backends/ed_diffusers.py @@ -3,6 +3,7 @@ from sdkit_common import ( stop_backend, install_backend, uninstall_backend, + is_installed, create_sdkit_context, ping, load_model, diff --git a/ui/easydiffusion/backends/sdkit_common.py b/ui/easydiffusion/backends/sdkit_common.py index d7a49c3e..4c7eba59 100644 --- a/ui/easydiffusion/backends/sdkit_common.py +++ b/ui/easydiffusion/backends/sdkit_common.py @@ -28,6 +28,10 @@ def uninstall_backend(): pass +def is_installed(): + return True + + def create_sdkit_context(use_diffusers): c = Context() c.test_diffusers = use_diffusers diff --git a/ui/easydiffusion/backends/webui/__init__.py b/ui/easydiffusion/backends/webui/__init__.py index 94d9ff03..95cd7883 100644 --- a/ui/easydiffusion/backends/webui/__init__.py +++ b/ui/easydiffusion/backends/webui/__init__.py @@ -4,6 +4,7 @@ import subprocess import threading from threading import local import psutil +import shutil from easydiffusion.app import ROOT_DIR, getConfig from easydiffusion.model_manager import get_model_dir @@ -29,10 +30,15 @@ ed_info = { "type": "backend", } +WEBUI_REPO = "https://github.com/lllyasviel/stable-diffusion-webui-forge.git" +WEBUI_COMMIT = "f4d5e8cac16a42fa939e78a0956b4c30e2b47bb5" + BACKEND_DIR = os.path.abspath(os.path.join(ROOT_DIR, "webui")) SYSTEM_DIR = os.path.join(BACKEND_DIR, "system") WEBUI_DIR = os.path.join(BACKEND_DIR, "webui") +OS_NAME = platform.system() + MODELS_TO_OVERRIDE = { "stable-diffusion": "--ckpt-dir", "vae": "--vae-dir", @@ -46,10 +52,42 @@ MODELS_TO_OVERRIDE = { } backend_process = None +conda = "conda" + + +def locate_conda(): + global conda + + which = "where" if OS_NAME == "Windows" else "which" + conda = subprocess.getoutput(f"{which} conda") + conda = conda.split("\n") + conda = conda[0].strip() + print("conda: ", conda) + + +locate_conda() def install_backend(): - pass + print("Installing the WebUI backend..") + + # create the conda env + run([conda, "create", "-y", "--prefix", SYSTEM_DIR], cwd=ROOT_DIR) + + # install python 3.10 and git in the conda env + run([conda, "install", "-y", "--prefix", SYSTEM_DIR, "-c", "conda-forge", "python=3.10", "git"], cwd=ROOT_DIR) + + # print info + run_in_conda(["git", "--version"], cwd=ROOT_DIR) + run_in_conda(["python", "--version"], cwd=ROOT_DIR) + + # clone webui + run_in_conda(["git", "clone", WEBUI_REPO, WEBUI_DIR], cwd=ROOT_DIR) + + # install cpu-only torch if the PC doesn't have a graphics card (for Windows and Linux). + # this avoids WebUI installing a CUDA version and trying to activate it + if OS_NAME in ("Windows", "Linux") and not has_discrete_graphics_card(): + run_in_conda(["python", "-m", "pip", "install", "torch", "torchvision"], cwd=WEBUI_DIR) def start_backend(): @@ -59,6 +97,18 @@ def start_backend(): if not os.path.exists(BACKEND_DIR): install_backend() + if backend_config.get("auto_update", True): + run_in_conda(["git", "add", "-A", "."], cwd=WEBUI_DIR) + run_in_conda(["git", "stash"], cwd=WEBUI_DIR) + run_in_conda(["git", "reset", "--hard"], cwd=WEBUI_DIR) + run_in_conda(["git", "fetch"], cwd=WEBUI_DIR) + run_in_conda(["git", "-c", "advice.detachedHead=false", "checkout", WEBUI_COMMIT], cwd=WEBUI_DIR) + + # hack to prevent webui-macos-env.sh from overwriting the COMMANDLINE_ARGS env variable + mac_webui_file = os.path.join(WEBUI_DIR, "webui-macos-env.sh") + if os.path.exists(mac_webui_file): + os.remove(mac_webui_file) + impl.WEBUI_HOST = backend_config.get("host", "localhost") impl.WEBUI_PORT = backend_config.get("port", "7860") @@ -68,11 +118,11 @@ def start_backend(): def target(): global backend_process - cmd = "webui.bat" if platform.system() == "Windows" else "webui.sh" + cmd = "webui.bat" if OS_NAME == "Windows" else "./webui.sh" while True: print("starting", cmd, WEBUI_DIR) - backend_process = subprocess.Popen([cmd], shell=True, cwd=WEBUI_DIR, env=env) + backend_process = run_in_conda([cmd], cwd=WEBUI_DIR, env=env, wait=False) backend_process.wait() stop_backend() @@ -94,7 +144,51 @@ def stop_backend(): def uninstall_backend(): - pass + shutil.rmtree(BACKEND_DIR) + + +def is_installed(): + if not os.path.exists(BACKEND_DIR) or not os.path.exists(SYSTEM_DIR) or not os.path.exists(WEBUI_DIR): + return True + + env = dict(os.environ) + env.update(get_env()) + + try: + out = check_output_in_conda(["python", "-m", "pip", "show", "torch"], env=env) + return "Version" in out.decode() + except subprocess.CalledProcessError: + pass + + return False + + +def run(cmds: list, cwd=None, env=None, stream_output=True, wait=True): + p = subprocess.Popen(cmds, cwd=cwd, env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + + if stream_output: + while True: + output = p.stdout.readline() + output = output.decode() + if output == "" and p.poll() is not None: + break + if output: + print(output, end="") + + if wait: + p.wait() + + return p + + +def run_in_conda(cmds: list, *args, **kwargs): + cmds = [conda, "run", "--no-capture-output", "--prefix", SYSTEM_DIR] + cmds + return run(cmds, *args, **kwargs) + + +def check_output_in_conda(cmds: list, cwd=None, env=None): + cmds = [conda, "run", "--no-capture-output", "--prefix", SYSTEM_DIR] + cmds + return subprocess.check_output(cmds, cwd=cwd, env=env, stderr=subprocess.PIPE) def create_context(): @@ -130,34 +224,62 @@ def get_env(): env_entries = { "PATH": [ - f"{dir}/git/bin", - f"{dir}/python", - f"{dir}/python/Library/bin", - f"{dir}/python/Scripts", - f"{dir}/python/Library/usr/bin", + f"{dir}", + f"{dir}/bin", + f"{dir}/Library/bin", + f"{dir}/Scripts", + f"{dir}/usr/bin", ], "PYTHONPATH": [ - f"{dir}/python", - f"{dir}/python/lib/site-packages", - f"{dir}/python/lib/python3.10/site-packages", + f"{dir}", + f"{dir}/lib/site-packages", + f"{dir}/lib/python3.10/site-packages", ], "PYTHONHOME": [], - "PY_LIBS": [f"{dir}/python/Scripts/Lib", f"{dir}/python/Scripts/Lib/site-packages"], - "PY_PIP": [f"{dir}/python/Scripts"], - "PIP_INSTALLER_LOCATION": [f"{dir}/python/get-pip.py"], + "PY_LIBS": [ + f"{dir}/Scripts/Lib", + f"{dir}/Scripts/Lib/site-packages", + f"{dir}/lib", + f"{dir}/lib/python3.10/site-packages", + ], + "PY_PIP": [f"{dir}/Scripts", f"{dir}/bin"], + "PIP_INSTALLER_LOCATION": [], # [f"{dir}/python/get-pip.py"], "TRANSFORMERS_CACHE": [f"{dir}/transformers-cache"], "HF_HUB_DISABLE_SYMLINKS_WARNING": ["true"], - "COMMANDLINE_ARGS": [f'--api --models-dir "{models_dir}" {model_path_args}'], + "COMMANDLINE_ARGS": [f'--api --models-dir "{models_dir}" {model_path_args} --skip-torch-cuda-test'], "SKIP_VENV": ["1"], "SD_WEBUI_RESTARTING": ["1"], - "PYTHON": [f"{dir}/python/python"], - "GIT": [f"{dir}/git/bin/git"], } - if platform.system() == "Windows": + if OS_NAME == "Windows": + env_entries["PATH"].append("C:/Windows/System32") + env_entries["PATH"].append("C:/Windows/System32/wbem") env_entries["PYTHONNOUSERSITE"] = ["1"] + env_entries["PYTHON"] = [f"{dir}/python"] + env_entries["GIT"] = [f"{dir}/Library/bin/git"] else: + env_entries["PATH"].append("/bin") + env_entries["PATH"].append("/usr/bin") + env_entries["PATH"].append("/usr/sbin") env_entries["PYTHONNOUSERSITE"] = ["y"] + env_entries["PYTHON"] = [f"{dir}/bin/python"] + env_entries["GIT"] = [f"{dir}/bin/git"] + env_entries["venv_dir"] = ["-"] + + if OS_NAME in ("Windows", "Linux") and not has_discrete_graphics_card(): + env_entries["COMMANDLINE_ARGS"][0] += " --always-cpu" + + if OS_NAME == "Darwin": + # based on https://github.com/lllyasviel/stable-diffusion-webui-forge/blob/e26abf87ecd1eefd9ab0a198eee56f9c643e4001/webui-macos-env.sh + # hack - have to define these here, otherwise webui-macos-env.sh will overwrite COMMANDLINE_ARGS + env_entries["COMMANDLINE_ARGS"][0] += " --upcast-sampling --no-half-vae --use-cpu interrogate" + env_entries["PYTORCH_ENABLE_MPS_FALLBACK"] = ["1"] + + cpu_name = str(subprocess.check_output(["sysctl", "-n", "machdep.cpu.brand_string"])) + if "Intel" in cpu_name: + env_entries["TORCH_COMMAND"] = ["pip install torch==2.1.2 torchvision==0.16.2"] + else: + env_entries["TORCH_COMMAND"] = ["pip install torch==2.3.1 torchvision==0.18.1"] env = {} for key, paths in env_entries.items(): @@ -169,6 +291,40 @@ def get_env(): return env +def has_discrete_graphics_card(): + system = OS_NAME + + if system == "Windows": + try: + output = subprocess.check_output( + ["wmic", "path", "win32_videocontroller", "get", "name"], stderr=subprocess.STDOUT + ) + # Filter for discrete graphics cards (NVIDIA, AMD, etc.) + discrete_gpus = ["NVIDIA", "AMD", "ATI"] + return any(gpu in output.decode() for gpu in discrete_gpus) + except subprocess.CalledProcessError: + return False + + elif system == "Linux": + try: + output = subprocess.check_output(["lspci"], stderr=subprocess.STDOUT) + # Check for discrete GPUs (NVIDIA, AMD) + discrete_gpus = ["NVIDIA", "AMD", "Advanced Micro Devices"] + return any(gpu in line for line in output.decode().splitlines() for gpu in discrete_gpus) + except subprocess.CalledProcessError: + return False + + elif system == "Darwin": # macOS + try: + output = subprocess.check_output(["system_profiler", "SPDisplaysDataType"], stderr=subprocess.STDOUT) + # Check for discrete GPU in the output + return "NVIDIA" in output.decode() or "AMD" in output.decode() + except subprocess.CalledProcessError: + return False + + return False + + # https://stackoverflow.com/a/25134985 def kill(proc_pid): process = psutil.Process(proc_pid) diff --git a/ui/easydiffusion/task_manager.py b/ui/easydiffusion/task_manager.py index 5ad6420d..38cba0b6 100644 --- a/ui/easydiffusion/task_manager.py +++ b/ui/easydiffusion/task_manager.py @@ -234,6 +234,7 @@ def thread_render(device): from easydiffusion import model_manager, runtime from easydiffusion.backend_manager import backend + from requests import ConnectionError try: runtime.init(device) @@ -252,7 +253,7 @@ def thread_render(device): break time.sleep(1) - except TimeoutError: + except (TimeoutError, ConnectionError): time.sleep(1) model_manager.load_default_models(runtime.context)