From 1f7c7909c201081da1923ab56887af88f67e4b31 Mon Sep 17 00:00:00 2001 From: cmdr2 Date: Mon, 10 Oct 2022 19:35:33 +0530 Subject: [PATCH] Initial port of the entire installation process; Switched 0.0.0.0 to localhost default; Skip color correction in GFPGAN via a patch --- Start Stable Diffusion UI.cmd | 1 - installer/installer/app.py | 40 ++++++++---- installer/installer/check_modules.py | 9 ++- installer/installer/helpers.py | 65 +++++++++---------- installer/installer/main.py | 6 ++ .../installer/tasks/apply_project_update.py | 4 ++ installer/installer/tasks/download_weights.py | 46 +++++++++++++ .../tasks/fetch_stable_diffusion_repo.py | 18 ++--- .../install_stable_diffusion_packages.py | 22 +++++-- .../installer/tasks/install_ui_packages.py | 34 ++++++++++ installer/installer/tasks/start_ui_server.py | 23 +++++++ installer/patches/gfpgan_custom.patch | 22 +++++++ installer/yaml/sd-environment-mac-nvidia.yaml | 2 + .../yaml/sd-environment-win-linux-nvidia.yaml | 2 + start.sh | 1 - 15 files changed, 224 insertions(+), 71 deletions(-) create mode 100644 installer/installer/tasks/download_weights.py create mode 100644 installer/installer/tasks/install_ui_packages.py create mode 100644 installer/installer/tasks/start_ui_server.py create mode 100644 installer/patches/gfpgan_custom.patch diff --git a/Start Stable Diffusion UI.cmd b/Start Stable Diffusion UI.cmd index e13dae88..8452cc54 100644 --- a/Start Stable Diffusion UI.cmd +++ b/Start Stable Diffusion UI.cmd @@ -4,7 +4,6 @@ echo. & echo "Stable Diffusion UI - v2.5" & echo. set PATH=C:\Windows\System32;%PATH% -set START_CMD_FILENAME=Start Stable Diffusion UI.cmd set SD_BASE_DIR=%cd% @rem Confirm or change the installation dir diff --git a/installer/installer/app.py b/installer/installer/app.py index 6e0493a9..762d137d 100644 --- a/installer/installer/app.py +++ b/installer/installer/app.py @@ -1,5 +1,6 @@ import os import json +import platform # config PROJECT_REPO_URL = 'https://github.com/cmdr2/stable-diffusion-ui.git' @@ -12,13 +13,14 @@ STABLE_DIFFUSION_REPO_DIR_NAME = 'stable-diffusion' PROJECT_ENV_DIR_NAME = 'project_env' -START_CMD_FILE_NAME = os.environ['START_CMD_FILENAME'] -LOG_FILE_NAME = 'run.log' +START_CMD_FILE_NAME = "Start Stable Diffusion UI.cmd" if platform.system == "Windows" else "start.sh" +DEV_CONSOLE_CMD_FILE_NAME = "Developer Console.cmd" if platform.system == "Windows" else "developer_console.sh" CONFIG_FILE_NAME = 'config.json' # top-level folders ENV_DIR_NAME = 'env' +MODELS_DIR_NAME = 'models' INSTALLER_DIR_NAME = 'installer' UI_DIR_NAME = 'ui' @@ -28,16 +30,13 @@ ENGINE_DIR_NAME = 'engine' # env SD_BASE_DIR = os.environ['SD_BASE_DIR'] -def get_config(): - config_path = os.path.join(SD_BASE_DIR, CONFIG_FILE_NAME) - if not os.path.exists(config_path): - return {} - with open(config_path, "r") as f: - return json.load(f) +# model folders +STABLE_DIFFUSION_MODELS_DIR_NAME = "stable-diffusion" +GFPGAN_MODELS_DIR_NAME = "gfpgan" +RealESRGAN_MODELS_DIR_NAME = "realesrgan" - -# references +# create references to dirs env_dir_path = os.path.join(SD_BASE_DIR, ENV_DIR_NAME) installer_dir_path = os.path.join(SD_BASE_DIR, INSTALLER_DIR_NAME) @@ -49,5 +48,24 @@ stable_diffusion_repo_dir_path = os.path.join(env_dir_path, STABLE_DIFFUSION_REP project_env_dir_path = os.path.join(env_dir_path, PROJECT_ENV_DIR_NAME) +patches_dir_path = os.path.join(installer_dir_path, 'patches') + +models_dir_path = os.path.join(SD_BASE_DIR, MODELS_DIR_NAME) +stable_diffusion_models_dir_path = os.path.join(models_dir_path, STABLE_DIFFUSION_MODELS_DIR_NAME) +gfpgan_models_dir_path = os.path.join(models_dir_path, GFPGAN_MODELS_DIR_NAME) +realesrgan_models_dir_path = os.path.join(models_dir_path, RealESRGAN_MODELS_DIR_NAME) + + +# useful functions +def get_config(): + config_path = os.path.join(SD_BASE_DIR, CONFIG_FILE_NAME) + if not os.path.exists(config_path): + return {} + + with open(config_path, "r") as f: + return json.load(f) + + +# app context config = get_config() -log_file = open(LOG_FILE_NAME, 'wb') +activated_env_dir_path = None diff --git a/installer/installer/check_modules.py b/installer/installer/check_modules.py index fb566085..296907bc 100644 --- a/installer/installer/check_modules.py +++ b/installer/installer/check_modules.py @@ -1,3 +1,7 @@ +''' +This script is run by the `installer.helpers.modules_exist_in_env()` function +''' + import sys import pkgutil @@ -9,7 +13,6 @@ for m in modules: if len(missing_modules) == 0: print('42') - exit(0) + exit() -print('Missing modules', missing_modules) -exit(1) \ No newline at end of file +print('Missing modules', missing_modules) \ No newline at end of file diff --git a/installer/installer/helpers.py b/installer/installer/helpers.py index 8fb49d90..31c73f8d 100644 --- a/installer/installer/helpers.py +++ b/installer/installer/helpers.py @@ -1,57 +1,46 @@ from os import path import subprocess -import sys -import shutil -import time from installer import app -def run(cmd, run_in_folder=None, get_output=False, write_to_log=True, env=None): +def run(cmd, run_in_folder=None, env=None, get_output=False, log_the_cmd=False): + if app.activated_env_dir_path is not None and 'micromamba activate' not in cmd: + cmd = f'micromamba activate "{app.activated_env_dir_path}" && {cmd}' + if run_in_folder is not None: cmd = f'cd "{run_in_folder}" && {cmd}' - p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=True, env=env) - - buf = bytearray() - - for c in iter(lambda: p.stdout.read(1), b""): - sys.stdout.buffer.write(c) - sys.stdout.flush() - - buf.extend(c) - - if write_to_log and app.log_file is not None: - app.log_file.write(c) - app.log_file.flush() - - p.wait() + if log_the_cmd: + log('running: ' + cmd) if get_output: - return p.returncode, buf.decode('utf-8') + p = subprocess.Popen(cmd, shell=True, env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + else: + p = subprocess.Popen(cmd, shell=True, env=env) - return p.returncode == 0 + out, err = p.communicate() + + if get_output: + return out, err def log(msg): print(msg) - app.log_file.write(bytes(msg + "\n", 'utf-8')) - app.log_file.flush() - def modules_exist_in_env(modules, env_dir_path=app.project_env_dir_path): if not path.exists(env_dir_path): return False - activate_cmd = f'micromamba activate "{env_dir_path}"' - - if not run(activate_cmd, write_to_log=False): - return False - check_modules_script_path = path.join(app.installer_dir_path, 'installer', 'check_modules.py') module_args = ' '.join(modules) - check_modules_cmd = f'{activate_cmd} && python "{check_modules_script_path}" {module_args}' + check_modules_cmd = f'python "{check_modules_script_path}" {module_args}' - ret_code, output = run(check_modules_cmd, get_output=True, write_to_log=False) - if ret_code != 0 or 'Missing' in output: + if app.activated_env_dir_path != env_dir_path: + activate_cmd = f'micromamba activate "{env_dir_path}"' + check_modules_cmd = f'{activate_cmd} && {check_modules_cmd}' + + # activate and run the modules checker + output, _ = run(check_modules_cmd, get_output=True) + if 'Missing' in output: return False return True @@ -66,10 +55,16 @@ Error: {error_msg}. Sorry about that, please try to: 3. If those steps don't help, please copy *all* the error messages in this window, and ask the community at https://discord.com/invite/u9yhsFmEkB 4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues Thanks!''') - - ts = int(time.time()) - shutil.copy(app.LOG_FILE_NAME, f'error-{ts}.log') except: pass exit(1) + +def apply_git_patches(repo_dir_path, patch_file_names): + is_developer_mode = app.config.get('is_developer_mode', False) + if is_developer_mode: + return + + for patch_file_name in patch_file_names: + patch_file_path = path.join(app.patches_dir_path, patch_file_name) + run(f"git apply {patch_file_path}", run_in_folder=repo_dir_path) diff --git a/installer/installer/main.py b/installer/installer/main.py index d3fcca7b..fa76e469 100644 --- a/installer/installer/main.py +++ b/installer/installer/main.py @@ -10,6 +10,9 @@ from installer.tasks import ( apply_project_update, fetch_stable_diffusion_repo, install_stable_diffusion_packages, + install_ui_packages, + download_weights, + start_ui_server, ) tasks = [ @@ -17,6 +20,9 @@ tasks = [ apply_project_update, fetch_stable_diffusion_repo, install_stable_diffusion_packages, + install_ui_packages, + download_weights, + start_ui_server, ] helpers.log(f'Starting Stable Diffusion UI at {datetime.now().strftime("%d/%m/%Y %H:%M:%S")}') diff --git a/installer/installer/tasks/apply_project_update.py b/installer/installer/tasks/apply_project_update.py index 5c930bf5..e57249a9 100644 --- a/installer/installer/tasks/apply_project_update.py +++ b/installer/installer/tasks/apply_project_update.py @@ -15,6 +15,9 @@ def run(): start_cmd_src_path = path.join(app.project_repo_dir_path, app.START_CMD_FILE_NAME) start_cmd_dst_path = path.join(app.SD_BASE_DIR, app.START_CMD_FILE_NAME) + dev_console_cmd_src_path = path.join(app.project_repo_dir_path, app.DEV_CONSOLE_CMD_FILE_NAME) + dev_console_cmd_dst_path = path.join(app.SD_BASE_DIR, app.DEV_CONSOLE_CMD_FILE_NAME) + shutil.rmtree(app.installer_dir_path, ignore_errors=True) shutil.rmtree(app.ui_dir_path, ignore_errors=True) shutil.rmtree(app.engine_dir_path, ignore_errors=True) @@ -24,3 +27,4 @@ def run(): shutil.copytree(engine_src_path, app.engine_dir_path, dirs_exist_ok=True) shutil.copy(start_cmd_src_path, start_cmd_dst_path) + shutil.copy(dev_console_cmd_src_path, dev_console_cmd_dst_path) diff --git a/installer/installer/tasks/download_weights.py b/installer/installer/tasks/download_weights.py new file mode 100644 index 00000000..b21ac159 --- /dev/null +++ b/installer/installer/tasks/download_weights.py @@ -0,0 +1,46 @@ +import os + +from installer import app, helpers + +def run(): + fetch_model('Stable Diffusion', 'sd-v1-4.ckpt', model_dir_path=app.stable_diffusion_models_dir_path, download_url='https://me.cmdr2.org/stable-diffusion-ui/sd-v1-4.ckpt', expected_file_sizes=[4265380512, 7703807346, 7703810927]) + fetch_model('Face Correction (GFPGAN)', 'GFPGANv1.4.pth', model_dir_path=app.gfpgan_models_dir_path, download_url='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.4.pth', expected_file_sizes=[348632874]) + fetch_model('Resolution Upscale (RealESRGAN x4)', 'RealESRGAN_x4plus.pth', model_dir_path=app.realesrgan_models_dir_path, download_url='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth', expected_file_sizes=[67040989]) + fetch_model('Resolution Upscale (RealESRGAN x4_anime)', 'RealESRGAN_x4plus_anime_6B.pth', model_dir_path=app.realesrgan_models_dir_path, download_url='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth', expected_file_sizes=[17938799]) + +def fetch_model(model_type, file_name, model_dir_path, download_url, expected_file_sizes): + os.makedirs(model_dir_path, exist_ok=True) + + file_path = os.path.join(model_dir_path, file_name) + + if model_exists(file_name, file_path, expected_file_sizes): + helpers.log(f'Data files (weights) necessary for {model_type} were already downloaded') + return + + helpers.log(f'Downloading data files (weights) for {model_type}..') + + helpers.run(f'curl -L -k "{download_url}" > "{file_path}"', log_the_cmd=True) + +def model_exists(file_name, file_path, expected_file_sizes): + legacy_file_path = os.path.join(app.stable_diffusion_repo_dir_path, file_name) + + file_exists = os.path.exists(file_path) + legacy_file_exists = os.path.exists(legacy_file_path) + + if legacy_file_exists: + file_size = os.path.getsize(legacy_file_path) + if file_size in expected_file_sizes: + return True + + helpers.log(f'{file_name} is invalid. Was only {file_size} bytes in size. Downloading again..') + os.remove(legacy_file_path) + + if file_exists: + file_size = os.path.getsize(file_path) + if file_size in expected_file_sizes: + return True + + helpers.log(f'{file_name} is invalid. Was only {file_size} bytes in size. Downloading again..') + os.remove(file_path) + + return False diff --git a/installer/installer/tasks/fetch_stable_diffusion_repo.py b/installer/installer/tasks/fetch_stable_diffusion_repo.py index cd92a0b8..2cd0a4c5 100644 --- a/installer/installer/tasks/fetch_stable_diffusion_repo.py +++ b/installer/installer/tasks/fetch_stable_diffusion_repo.py @@ -2,18 +2,16 @@ from os import path from installer import app, helpers -patch_file_names = [ - 'sd_custom.patch', -] - stable_diffusion_repo_git_path = path.join(app.stable_diffusion_repo_dir_path, '.git') -patches_dir_path = path.join(app.installer_dir_path, 'patches') is_developer_mode = app.config.get('is_developer_mode', False) def run(): fetch_repo() - apply_patches() + + helpers.apply_git_patches(app.stable_diffusion_repo_dir_path, patch_file_names=( + "sd_custom.patch", + )) def fetch_repo(): commit_id = app.config.get('stable_diffusion_commit', app.DEFAULT_STABLE_DIFFUSION_COMMIT) @@ -37,11 +35,3 @@ def fetch_repo(): helpers.fail_with_install_error(error_msg="Could not download Stable Diffusion") helpers.run(f'git -c advice.detachedHead=false checkout "{commit_id}"', run_in_folder=app.stable_diffusion_repo_dir_path) - -def apply_patches(): - if is_developer_mode: - return - - for patch_file_name in patch_file_names: - patch_file_path = path.join(patches_dir_path, patch_file_name) - helpers.run(f"git apply {patch_file_path}", run_in_folder=app.stable_diffusion_repo_dir_path) diff --git a/installer/installer/tasks/install_stable_diffusion_packages.py b/installer/installer/tasks/install_stable_diffusion_packages.py index d9d134bd..7574d876 100644 --- a/installer/installer/tasks/install_stable_diffusion_packages.py +++ b/installer/installer/tasks/install_stable_diffusion_packages.py @@ -1,5 +1,4 @@ import os -import shutil import platform from installer import app, helpers @@ -11,20 +10,31 @@ def run(): log_installing_header() - shutil.rmtree(app.project_env_dir_path, ignore_errors=True) - environment_file_path = get_environment_file_path() env = os.environ.copy() env['PYTHONNOUSERSITE'] = '1' - if helpers.run(f'micromamba create --prefix {app.project_env_dir_path} -f {environment_file_path}', env=env) \ - and is_valid_env(): + if not os.path.exists(app.project_env_dir_path): + helpers.run(f'micromamba create --prefix {app.project_env_dir_path}', log_the_cmd=True) + helpers.run(f'micromamba install -y --prefix {app.project_env_dir_path} -f {environment_file_path}', env=env, log_the_cmd=True) + + if is_valid_env(): helpers.log("Installed the packages necessary for Stable Diffusion") + + app.activated_env_dir_path = app.project_env_dir_path # so that future `run()` invocations will run in the activated env else: helpers.fail_with_install_error(error_msg="Could not install the packages necessary for Stable Diffusion") + apply_patches() + +def apply_patches(): + gfpgan_repo_dir_path = os.path.join(app.stable_diffusion_repo_dir_path, 'src', 'gfpgan') + helpers.apply_git_patches(gfpgan_repo_dir_path, patch_file_names=( + "gfpgan_custom.patch", + )) + def get_environment_file_path(): environment_file_name = 'sd-environment-win-linux-nvidia.yaml' if platform.system() == 'Darwin': @@ -42,4 +52,4 @@ Downloading packages necessary for Stable Diffusion.. ''') def is_valid_env(): - return helpers.modules_exist_in_env(('torch', 'ldm', 'antlr4', 'transformers', 'numpy')) + return helpers.modules_exist_in_env(('torch', 'ldm', 'antlr4', 'transformers', 'numpy', 'gfpgan', 'realesrgan', 'basicsr')) diff --git a/installer/installer/tasks/install_ui_packages.py b/installer/installer/tasks/install_ui_packages.py new file mode 100644 index 00000000..bfbd3854 --- /dev/null +++ b/installer/installer/tasks/install_ui_packages.py @@ -0,0 +1,34 @@ +import os +import shutil + +from installer import app, helpers + +def run(): + if is_valid_env(): + helpers.log("Packages necessary for Stable Diffusion UI were already installed") + return + + log_installing_header() + + env = os.environ.copy() + env['PYTHONNOUSERSITE'] = '1' + + helpers.run(f'micromamba install -y --prefix {app.project_env_dir_path} -c conda-forge uvicorn fastapi', env=env, log_the_cmd=True) + + if is_valid_env(): + helpers.log("Installed the packages necessary for Stable Diffusion UI") + else: + helpers.fail_with_install_error(error_msg="Could not install the packages necessary for Stable Diffusion UI") + +def log_installing_header(): + helpers.log(''' + +Downloading packages necessary for Stable Diffusion UI.. + +''') + +def is_valid_env(): + if shutil.which("uvicorn") is None: + return False + + return helpers.modules_exist_in_env(('uvicorn', 'fastapi')) diff --git a/installer/installer/tasks/start_ui_server.py b/installer/installer/tasks/start_ui_server.py new file mode 100644 index 00000000..424cd5ed --- /dev/null +++ b/installer/installer/tasks/start_ui_server.py @@ -0,0 +1,23 @@ +import os +import shutil +import platform + +from installer import app, helpers + +def run(): + helpers.log("\nStable Diffusion is ready!\n") + + env = os.environ.copy() + env['SD_DIR'] = app.stable_diffusion_repo_dir_path + env['PYTHONPATH'] = app.stable_diffusion_repo_dir_path + ';' + os.path.join(app.project_env_dir_path, 'lib', 'site-packages') + env['SD_UI_PATH'] = app.ui_dir_path + + helpers.log(f'PYTHONPATH={env["PYTHONPATH"]}') + helpers.run('python --version', log_the_cmd=True) + + host = app.config.get('host', 'localhost') + port = app.config.get('port', '9000') + + ui_server_cmd = f'uvicorn server:app --app-dir "{app.ui_dir_path}" --port {port} --host {host}' + + helpers.run(ui_server_cmd, run_in_folder=app.stable_diffusion_repo_dir_path, log_the_cmd=True, env=env) diff --git a/installer/patches/gfpgan_custom.patch b/installer/patches/gfpgan_custom.patch new file mode 100644 index 00000000..2213275c --- /dev/null +++ b/installer/patches/gfpgan_custom.patch @@ -0,0 +1,22 @@ +diff --git a/gfpgan/utils.py b/gfpgan/utils.py +index 74ee5a8..1357f48 100644 +--- a/gfpgan/utils.py ++++ b/gfpgan/utils.py +@@ -117,14 +117,14 @@ class GFPGANer(): + # face restoration + for cropped_face in self.face_helper.cropped_faces: + # prepare data +- cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True) ++ cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=False, float32=True) + normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True) + cropped_face_t = cropped_face_t.unsqueeze(0).to(self.device) + + try: +- output = self.gfpgan(cropped_face_t, return_rgb=False, weight=weight)[0] ++ output = self.gfpgan(cropped_face_t, return_rgb=True, weight=weight)[0] + # convert to image +- restored_face = tensor2img(output.squeeze(0), rgb2bgr=True, min_max=(-1, 1)) ++ restored_face = tensor2img(output.squeeze(0), rgb2bgr=False, min_max=(-1, 1)) + except RuntimeError as error: + print(f'\tFailed inference for GFPGAN: {error}.') + restored_face = cropped_face diff --git a/installer/yaml/sd-environment-mac-nvidia.yaml b/installer/yaml/sd-environment-mac-nvidia.yaml index 80746ffd..8f2a4c77 100644 --- a/installer/yaml/sd-environment-mac-nvidia.yaml +++ b/installer/yaml/sd-environment-mac-nvidia.yaml @@ -40,6 +40,8 @@ dependencies: - torch-fidelity==0.3.0 - -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers - -e git+https://github.com/openai/CLIP.git@main#egg=clip + - -e git+https://github.com/TencentARC/GFPGAN#egg=GFPGAN + - -e git+https://github.com/xinntao/Real-ESRGAN#egg=realesrgan - -e . variables: PYTORCH_ENABLE_MPS_FALLBACK: 1 diff --git a/installer/yaml/sd-environment-win-linux-nvidia.yaml b/installer/yaml/sd-environment-win-linux-nvidia.yaml index d92c0d9a..e2819d6b 100644 --- a/installer/yaml/sd-environment-win-linux-nvidia.yaml +++ b/installer/yaml/sd-environment-win-linux-nvidia.yaml @@ -28,4 +28,6 @@ dependencies: - kornia==0.6 - -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers - -e git+https://github.com/openai/CLIP.git@main#egg=clip + - -e git+https://github.com/TencentARC/GFPGAN#egg=GFPGAN + - -e git+https://github.com/xinntao/Real-ESRGAN#egg=realesrgan - -e . diff --git a/start.sh b/start.sh index 008abda4..78fcd5a3 100755 --- a/start.sh +++ b/start.sh @@ -3,7 +3,6 @@ echo "Stable Diffusion UI - v2.5" echo "" -export START_CMD_FILENAME="start.sh" export SD_BASE_DIR=$(pwd) echo "Working in $SD_BASE_DIR"