diff --git a/installer/installer/app.py b/installer/installer/app.py index 0bbfd6bf..6e0493a9 100644 --- a/installer/installer/app.py +++ b/installer/installer/app.py @@ -10,6 +10,8 @@ STABLE_DIFFUSION_REPO_URL = 'https://github.com/basujindal/stable-diffusion.git' DEFAULT_STABLE_DIFFUSION_COMMIT = 'f6cfebffa752ee11a7b07497b8529d5971de916c' 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' CONFIG_FILE_NAME = 'config.json' @@ -45,5 +47,7 @@ engine_dir_path = os.path.join(SD_BASE_DIR, ENGINE_DIR_NAME) project_repo_dir_path = os.path.join(env_dir_path, PROJECT_REPO_DIR_NAME) stable_diffusion_repo_dir_path = os.path.join(env_dir_path, STABLE_DIFFUSION_REPO_DIR_NAME) +project_env_dir_path = os.path.join(env_dir_path, PROJECT_ENV_DIR_NAME) + config = get_config() log_file = open(LOG_FILE_NAME, 'wb') diff --git a/installer/installer/check_modules.py b/installer/installer/check_modules.py new file mode 100644 index 00000000..fb566085 --- /dev/null +++ b/installer/installer/check_modules.py @@ -0,0 +1,15 @@ +import sys +import pkgutil + +modules = sys.argv[1:] +missing_modules = [] +for m in modules: + if pkgutil.find_loader(m) is None: + missing_modules.append(m) + +if len(missing_modules) == 0: + print('42') + exit(0) + +print('Missing modules', missing_modules) +exit(1) \ No newline at end of file diff --git a/installer/installer/helpers.py b/installer/installer/helpers.py index b64897d0..8fb49d90 100644 --- a/installer/installer/helpers.py +++ b/installer/installer/helpers.py @@ -1,3 +1,4 @@ +from os import path import subprocess import sys import shutil @@ -5,22 +6,29 @@ import time from installer import app -def run(cmd, run_in_folder=None): +def run(cmd, run_in_folder=None, get_output=False, write_to_log=True, env=None): 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) + 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() - if app.log_file is not None: + 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 get_output: + return p.returncode, buf.decode('utf-8') + return p.returncode == 0 def log(msg): @@ -29,6 +37,25 @@ def log(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}' + + ret_code, output = run(check_modules_cmd, get_output=True, write_to_log=False) + if ret_code != 0 or 'Missing' in output: + return False + + return True + def fail_with_install_error(error_msg): try: log(f''' diff --git a/installer/installer/main.py b/installer/installer/main.py index 7aaed91e..d3fcca7b 100644 --- a/installer/installer/main.py +++ b/installer/installer/main.py @@ -9,12 +9,14 @@ from installer.tasks import ( fetch_project_repo, apply_project_update, fetch_stable_diffusion_repo, + install_stable_diffusion_packages, ) tasks = [ fetch_project_repo, apply_project_update, fetch_stable_diffusion_repo, + install_stable_diffusion_packages, ] helpers.log(f'Starting Stable Diffusion UI at {datetime.now().strftime("%d/%m/%Y %H:%M:%S")}') diff --git a/installer/installer/tasks/install_stable_diffusion_packages.py b/installer/installer/tasks/install_stable_diffusion_packages.py new file mode 100644 index 00000000..d9d134bd --- /dev/null +++ b/installer/installer/tasks/install_stable_diffusion_packages.py @@ -0,0 +1,45 @@ +import os +import shutil +import platform + +from installer import app, helpers + +def run(): + if is_valid_env(): + helpers.log("Packages necessary for Stable Diffusion were already installed") + return + + 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(): + + helpers.log("Installed the packages necessary for Stable Diffusion") + else: + helpers.fail_with_install_error(error_msg="Could not install the packages necessary for Stable Diffusion") + +def get_environment_file_path(): + environment_file_name = 'sd-environment-win-linux-nvidia.yaml' + if platform.system() == 'Darwin': + environment_file_name = 'sd-environment-mac-nvidia.yaml' + + return os.path.join(app.installer_dir_path, 'yaml', environment_file_name) + +def log_installing_header(): + helpers.log(''' + +Downloading packages necessary for Stable Diffusion.. + +***** !! This will take some time (depending on the speed of the Internet connection) and may appear to be stuck, but please be patient ***** + +''') + +def is_valid_env(): + return helpers.modules_exist_in_env(('torch', 'ldm', 'antlr4', 'transformers', 'numpy')) diff --git a/installer/yaml/installer-environment.yaml b/installer/yaml/installer-environment.yaml index 7efe7c07..a1641c8e 100644 --- a/installer/yaml/installer-environment.yaml +++ b/installer/yaml/installer-environment.yaml @@ -4,4 +4,4 @@ channels: - conda-forge dependencies: - git - - python=3.8.13 + - python=3.10.5 diff --git a/installer/yaml/sd-environment-mac-nvidia.yaml b/installer/yaml/sd-environment-mac-nvidia.yaml new file mode 100644 index 00000000..80746ffd --- /dev/null +++ b/installer/yaml/sd-environment-mac-nvidia.yaml @@ -0,0 +1,45 @@ +name: ldm +channels: + - pytorch + - conda-forge +dependencies: + - python==3.10.5 + - pip==22.2.2 + + - pytorch + - torchvision + + - albumentations==1.2.1 + - coloredlogs==15.0.1 + - einops==0.4.1 + - grpcio==1.46.4 + - humanfriendly==10.0 + - imageio==2.21.2 + - imageio-ffmpeg==0.4.7 + - imgaug==0.4.0 + - kornia==0.6.7 + - mpmath==1.2.1 + - nomkl + - numpy==1.23.2 + - omegaconf==2.1.1 + - onnx==1.12.0 + - onnxruntime==1.12.1 + - pudb==2022.1 + - pytorch-lightning==1.6.5 + - scipy==1.9.1 + - streamlit==1.12.2 + - sympy==1.10.1 + - tensorboard==2.9.0 + - torchmetrics==0.9.3 + - antlr4-python3-runtime=4.8 + - pip: + - opencv-python==4.6.0 + - realesrgan==0.2.5.0 + - test-tube==0.7.5 + - transformers==4.21.2 + - 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 . +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 new file mode 100644 index 00000000..d9775f42 --- /dev/null +++ b/installer/yaml/sd-environment-win-linux-nvidia.yaml @@ -0,0 +1,32 @@ +name: ldm +channels: + - pytorch + - defaults +dependencies: + - python=3.10.5 + - pip=20.3 + - cudatoolkit=11.3 + - pytorch=1.11.0 + - torchvision=0.12.0 + - numpy=1.19.2 + - antlr4-python3-runtime=4.8 + - pip: + - albumentations==0.4.3 + - opencv-python==4.1.2.30 + - pudb==2019.2 + - imageio==2.9.0 + - imageio-ffmpeg==0.4.2 + - pytorch-lightning==1.4.2 + - omegaconf==2.1.1 + - test-tube>=0.7.5 + - streamlit>=0.73.1 + - einops==0.3.0 + - torch-fidelity==0.3.0 + - transformers==4.19.2 + - torchmetrics==0.6.0 + - pywavelets==1.3.0 + - pandas==1.4.4 + - 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 .