From 0b1968c0178b919fb8b24a539b7581ca545dcb97 Mon Sep 17 00:00:00 2001
From: cmdr2 <secondary.cmdr2@gmail.com>
Date: Tue, 4 Oct 2022 19:21:14 +0530
Subject: [PATCH] Task to install Stable Diffusion's environment

---
 installer/installer/app.py                    |  4 ++
 installer/installer/check_modules.py          | 15 +++++++
 installer/installer/helpers.py                | 33 ++++++++++++--
 installer/installer/main.py                   |  2 +
 .../install_stable_diffusion_packages.py      | 45 +++++++++++++++++++
 installer/yaml/installer-environment.yaml     |  2 +-
 installer/yaml/sd-environment-mac-nvidia.yaml | 45 +++++++++++++++++++
 .../yaml/sd-environment-win-linux-nvidia.yaml | 32 +++++++++++++
 8 files changed, 174 insertions(+), 4 deletions(-)
 create mode 100644 installer/installer/check_modules.py
 create mode 100644 installer/installer/tasks/install_stable_diffusion_packages.py
 create mode 100644 installer/yaml/sd-environment-mac-nvidia.yaml
 create mode 100644 installer/yaml/sd-environment-win-linux-nvidia.yaml

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 .