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

@@ -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

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

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)
+print('Missing modules', missing_modules)

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)

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")}')

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)

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

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)

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'))

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'))

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)

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

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

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 .

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"