fix ROCm version

- installer now installs the pretty dialog-based console launcher
- added dialogrc for custom colors
- add updater to download new launcher when users do an update

README.md

@@ -145,7 +145,7 @@ not supported.
     _For Linux with an AMD GPU:_
 
     ```sh
-    pip install InvokeAI --use-pep517 --extra-index-url https://download.pytorch.org/whl/rocm5.2
+    pip install InvokeAI --use-pep517 --extra-index-url https://download.pytorch.org/whl/rocm5.4.2
     ```
 
     _For Macintoshes, either Intel or M1/M2:_

binary_installer/install.bat.in

@@ -1,164 +0,0 @@
-@echo off
-
-@rem This script will install git (if not found on the PATH variable)
-@rem  using micromamba (an 8mb static-linked single-file binary, conda replacement).
-@rem For users who already have git, this step will be skipped.
-
-@rem Next, it'll download the project's source code.
-@rem Then it will download a self-contained, standalone Python and unpack it.
-@rem Finally, it'll create the Python virtual environment and preload the models.
-
-@rem This enables a user to install this project without manually installing git or Python
-
-@rem change to the script's directory
-PUSHD "%~dp0"
-
-set "no_cache_dir=--no-cache-dir"
-if "%1" == "use-cache" (
-    set "no_cache_dir="
-)
-
-echo ***** Installing InvokeAI.. *****
-@rem Config
-set INSTALL_ENV_DIR=%cd%\installer_files\env
-@rem https://mamba.readthedocs.io/en/latest/installation.html
-set MICROMAMBA_DOWNLOAD_URL=https://github.com/cmdr2/stable-diffusion-ui/releases/download/v1.1/micromamba.exe
-set RELEASE_URL=https://github.com/invoke-ai/InvokeAI
-set RELEASE_SOURCEBALL=/archive/refs/heads/main.tar.gz
-set PYTHON_BUILD_STANDALONE_URL=https://github.com/indygreg/python-build-standalone/releases/download
-set PYTHON_BUILD_STANDALONE=20221002/cpython-3.10.7+20221002-x86_64-pc-windows-msvc-shared-install_only.tar.gz
-
-set PACKAGES_TO_INSTALL=
-
-call git --version >.tmp1 2>.tmp2
-if "%ERRORLEVEL%" NEQ "0" set PACKAGES_TO_INSTALL=%PACKAGES_TO_INSTALL% git
-
-@rem Cleanup
-del /q .tmp1 .tmp2
-
-@rem (if necessary) install git into a contained environment
-if "%PACKAGES_TO_INSTALL%" NEQ "" (
-    @rem download micromamba
-    echo ***** Downloading micromamba from %MICROMAMBA_DOWNLOAD_URL% to micromamba.exe *****
-
-    call curl -L "%MICROMAMBA_DOWNLOAD_URL%" > micromamba.exe
-
-    @rem test the mamba binary
-    echo ***** Micromamba version:  *****
-    call micromamba.exe --version
-
-    @rem create the installer env
-    if not exist "%INSTALL_ENV_DIR%" (
-        call micromamba.exe create -y --prefix "%INSTALL_ENV_DIR%"
-    )
-
-    echo ***** Packages to install:%PACKAGES_TO_INSTALL% *****
-
-    call micromamba.exe install -y --prefix "%INSTALL_ENV_DIR%" -c conda-forge %PACKAGES_TO_INSTALL%
-
-    if not exist "%INSTALL_ENV_DIR%" (
-        echo ----- There was a problem while installing "%PACKAGES_TO_INSTALL%" using micromamba. Cannot continue. -----
-        pause
-        exit /b
-    )
-)
-
-del /q micromamba.exe
-
-@rem For 'git' only
-set PATH=%INSTALL_ENV_DIR%\Library\bin;%PATH%
-
-@rem Download/unpack/clean up InvokeAI release sourceball
-set err_msg=----- InvokeAI source download failed -----
-echo Trying to download "%RELEASE_URL%%RELEASE_SOURCEBALL%"
-curl -L %RELEASE_URL%%RELEASE_SOURCEBALL% --output InvokeAI.tgz
-if %errorlevel% neq 0 goto err_exit
-
-set err_msg=----- InvokeAI source unpack failed -----
-tar -zxf InvokeAI.tgz
-if %errorlevel% neq 0 goto err_exit
-
-del /q InvokeAI.tgz
-
-set err_msg=----- InvokeAI source copy failed -----
-cd InvokeAI-*
-xcopy . .. /e /h
-if %errorlevel% neq 0 goto err_exit
-cd ..
-
-@rem cleanup
-for /f %%i in ('dir /b InvokeAI-*') do rd /s /q %%i
-rd /s /q .dev_scripts .github docker-build tests
-del /q requirements.in requirements-mkdocs.txt shell.nix
-
-echo ***** Unpacked InvokeAI source *****
-
-@rem Download/unpack/clean up python-build-standalone
-set err_msg=----- Python download failed -----
-curl -L %PYTHON_BUILD_STANDALONE_URL%/%PYTHON_BUILD_STANDALONE% --output python.tgz
-if %errorlevel% neq 0 goto err_exit
-
-set err_msg=----- Python unpack failed -----
-tar -zxf python.tgz
-if %errorlevel% neq 0 goto err_exit
-
-del /q python.tgz
-
-echo ***** Unpacked python-build-standalone *****
-
-@rem create venv
-set err_msg=----- problem creating venv -----
-.\python\python -E -s -m venv .venv
-if %errorlevel% neq 0 goto err_exit
-call .venv\Scripts\activate.bat
-
-echo ***** Created Python virtual environment *****
-
-@rem Print venv's Python version
-set err_msg=----- problem calling venv's python -----
-echo We're running under
-.venv\Scripts\python --version
-if %errorlevel% neq 0 goto err_exit
-
-set err_msg=----- pip update failed -----
-.venv\Scripts\python -m pip install %no_cache_dir% --no-warn-script-location --upgrade pip wheel
-if %errorlevel% neq 0 goto err_exit
-
-echo ***** Updated pip and wheel *****
-
-set err_msg=----- requirements file copy failed -----
-copy binary_installer\py3.10-windows-x86_64-cuda-reqs.txt requirements.txt
-if %errorlevel% neq 0 goto err_exit
-
-set err_msg=----- main pip install failed -----
-.venv\Scripts\python -m pip install %no_cache_dir% --no-warn-script-location -r requirements.txt
-if %errorlevel% neq 0 goto err_exit
-
-echo ***** Installed Python dependencies *****
-
-set err_msg=----- InvokeAI setup failed -----
-.venv\Scripts\python -m pip install %no_cache_dir% --no-warn-script-location -e .
-if %errorlevel% neq 0 goto err_exit
-
-copy binary_installer\invoke.bat.in .\invoke.bat
-echo ***** Installed invoke launcher script ******
-
-@rem more cleanup
-rd /s /q binary_installer installer_files
-
-@rem preload the models
-call .venv\Scripts\python ldm\invoke\config\invokeai_configure.py
-set err_msg=----- model download clone failed -----
-if %errorlevel% neq 0 goto err_exit
-deactivate
-
-echo ***** Finished downloading models *****
-
-echo All done! Execute the file invoke.bat in this directory to start InvokeAI
-pause
-exit
-
-:err_exit
-    echo %err_msg%
-    pause
-    exit

binary_installer/install.sh.in

@@ -1,235 +0,0 @@
-#!/usr/bin/env bash
-
-# ensure we're in the correct folder in case user's CWD is somewhere else
-scriptdir=$(dirname "$0")
-cd "$scriptdir"
-
-set -euo pipefail
-IFS=$'\n\t'
-
-function _err_exit {
-    if test "$1" -ne 0
-    then
-        echo -e "Error code $1; Error caught was '$2'"
-        read -p "Press any key to exit..."
-        exit
-    fi
-}
-
-# This script will install git (if not found on the PATH variable)
-#  using micromamba (an 8mb static-linked single-file binary, conda replacement).
-# For users who already have git, this step will be skipped.
-
-# Next, it'll download the project's source code.
-# Then it will download a self-contained, standalone Python and unpack it.
-# Finally, it'll create the Python virtual environment and preload the models.
-
-# This enables a user to install this project without manually installing git or Python
-
-echo -e "\n***** Installing InvokeAI into $(pwd)... *****\n"
-
-export no_cache_dir="--no-cache-dir"
-if [ $# -ge 1 ]; then
-    if [ "$1" = "use-cache" ]; then
-        export no_cache_dir=""
-    fi
-fi
-
-
-OS_NAME=$(uname -s)
-case "${OS_NAME}" in
-    Linux*)     OS_NAME="linux";;
-    Darwin*)    OS_NAME="darwin";;
-    *)          echo -e "\n----- Unknown OS: $OS_NAME! This script runs only on Linux or macOS -----\n" && exit
-esac
-
-OS_ARCH=$(uname -m)
-case "${OS_ARCH}" in
-    x86_64*)    ;;
-    arm64*)     ;;
-    *)          echo -e "\n----- Unknown system architecture: $OS_ARCH! This script runs only on x86_64 or arm64 -----\n" && exit
-esac
-
-# https://mamba.readthedocs.io/en/latest/installation.html
-MAMBA_OS_NAME=$OS_NAME
-MAMBA_ARCH=$OS_ARCH
-if [ "$OS_NAME" == "darwin" ]; then
-    MAMBA_OS_NAME="osx"
-fi
-
-if [ "$OS_ARCH" == "linux" ]; then
-    MAMBA_ARCH="aarch64"
-fi
-
-if [ "$OS_ARCH" == "x86_64" ]; then
-    MAMBA_ARCH="64"
-fi
-
-PY_ARCH=$OS_ARCH
-if [ "$OS_ARCH" == "arm64" ]; then
-    PY_ARCH="aarch64"
-fi
-
-# Compute device ('cd' segment of reqs files) detect goes here
-# This needs a ton of work
-# Suggestions:
-#   - lspci
-#   - check $PATH for nvidia-smi, gtt CUDA/GPU version from output
-#   - Surely there's a similar utility for AMD?
-CD="cuda"
-if [ "$OS_NAME" == "darwin" ] && [ "$OS_ARCH" == "arm64" ]; then
-    CD="mps"
-fi
-
-# config
-INSTALL_ENV_DIR="$(pwd)/installer_files/env"
-MICROMAMBA_DOWNLOAD_URL="https://micro.mamba.pm/api/micromamba/${MAMBA_OS_NAME}-${MAMBA_ARCH}/latest"
-RELEASE_URL=https://github.com/invoke-ai/InvokeAI
-RELEASE_SOURCEBALL=/archive/refs/heads/main.tar.gz
-PYTHON_BUILD_STANDALONE_URL=https://github.com/indygreg/python-build-standalone/releases/download
-if [ "$OS_NAME" == "darwin" ]; then
-    PYTHON_BUILD_STANDALONE=20221002/cpython-3.10.7+20221002-${PY_ARCH}-apple-darwin-install_only.tar.gz
-elif [ "$OS_NAME" == "linux" ]; then
-    PYTHON_BUILD_STANDALONE=20221002/cpython-3.10.7+20221002-${PY_ARCH}-unknown-linux-gnu-install_only.tar.gz
-fi
-echo "INSTALLING $RELEASE_SOURCEBALL FROM $RELEASE_URL"
-
-PACKAGES_TO_INSTALL=""
-
-if ! hash "git" &>/dev/null; then PACKAGES_TO_INSTALL="$PACKAGES_TO_INSTALL git"; fi
-
-# (if necessary) install git and conda into a contained environment
-if [ "$PACKAGES_TO_INSTALL" != "" ]; then
-    # download micromamba
-    echo -e "\n***** Downloading micromamba from $MICROMAMBA_DOWNLOAD_URL to micromamba *****\n"
-
-    curl -L "$MICROMAMBA_DOWNLOAD_URL" | tar -xvjO bin/micromamba > micromamba
-
-    chmod u+x ./micromamba
-
-    # test the mamba binary
-    echo -e "\n***** Micromamba version: *****\n"
-    ./micromamba --version
-
-    # create the installer env
-    if [ ! -e "$INSTALL_ENV_DIR" ]; then
-        ./micromamba create -y --prefix "$INSTALL_ENV_DIR"
-    fi
-
-    echo -e "\n***** Packages to install:$PACKAGES_TO_INSTALL *****\n"
-
-    ./micromamba install -y --prefix "$INSTALL_ENV_DIR" -c conda-forge "$PACKAGES_TO_INSTALL"
-
-    if [ ! -e "$INSTALL_ENV_DIR" ]; then
-        echo -e "\n----- There was a problem while initializing micromamba. Cannot continue. -----\n"
-        exit
-    fi
-fi
-
-rm -f micromamba.exe
-
-export PATH="$INSTALL_ENV_DIR/bin:$PATH"
-
-# Download/unpack/clean up InvokeAI release sourceball
-_err_msg="\n----- InvokeAI source download failed -----\n"
-curl -L $RELEASE_URL/$RELEASE_SOURCEBALL --output InvokeAI.tgz
-_err_exit $? _err_msg
-_err_msg="\n----- InvokeAI source unpack failed -----\n"
-tar -zxf InvokeAI.tgz
-_err_exit $? _err_msg
-
-rm -f InvokeAI.tgz
-
-_err_msg="\n----- InvokeAI source copy failed -----\n"
-cd InvokeAI-*
-cp -r . ..
-_err_exit $? _err_msg
-cd ..
-
-# cleanup
-rm -rf InvokeAI-*/
-rm -rf .dev_scripts/ .github/ docker-build/ tests/ requirements.in requirements-mkdocs.txt shell.nix
-
-echo -e "\n***** Unpacked InvokeAI source *****\n"
-
-# Download/unpack/clean up python-build-standalone
-_err_msg="\n----- Python download failed -----\n"
-curl -L $PYTHON_BUILD_STANDALONE_URL/$PYTHON_BUILD_STANDALONE --output python.tgz
-_err_exit $? _err_msg
-_err_msg="\n----- Python unpack failed -----\n"
-tar -zxf python.tgz
-_err_exit $? _err_msg
-
-rm -f python.tgz
-
-echo -e "\n***** Unpacked python-build-standalone *****\n"
-
-# create venv
-_err_msg="\n----- problem creating venv -----\n"
-
-if [ "$OS_NAME" == "darwin" ]; then
-    # patch sysconfig so that extensions can build properly
-    # adapted from https://github.com/cashapp/hermit-packages/commit/fcba384663892f4d9cfb35e8639ff7a28166ee43
-    PYTHON_INSTALL_DIR="$(pwd)/python"
-    SYSCONFIG="$(echo python/lib/python*/_sysconfigdata_*.py)"
-    TMPFILE="$(mktemp)"
-    chmod +w "${SYSCONFIG}"
-    cp "${SYSCONFIG}" "${TMPFILE}"
-    sed "s,'/install,'${PYTHON_INSTALL_DIR},g" "${TMPFILE}" > "${SYSCONFIG}"
-    rm -f "${TMPFILE}"
-fi
-
-./python/bin/python3 -E -s -m venv .venv
-_err_exit $? _err_msg
-source .venv/bin/activate
-
-echo -e "\n***** Created Python virtual environment *****\n"
-
-# Print venv's Python version
-_err_msg="\n----- problem calling venv's python -----\n"
-echo -e "We're running under"
-.venv/bin/python3 --version
-_err_exit $? _err_msg
-
-_err_msg="\n----- pip update failed -----\n"
-.venv/bin/python3 -m pip install $no_cache_dir --no-warn-script-location --upgrade pip
-_err_exit $? _err_msg
-
-echo -e "\n***** Updated pip *****\n"
-
-_err_msg="\n----- requirements file copy failed -----\n"
-cp binary_installer/py3.10-${OS_NAME}-"${OS_ARCH}"-${CD}-reqs.txt requirements.txt
-_err_exit $? _err_msg
-
-_err_msg="\n----- main pip install failed -----\n"
-.venv/bin/python3 -m pip install $no_cache_dir --no-warn-script-location -r requirements.txt
-_err_exit $? _err_msg
-
-echo -e "\n***** Installed Python dependencies *****\n"
-
-_err_msg="\n----- InvokeAI setup failed -----\n"
-.venv/bin/python3 -m pip install $no_cache_dir --no-warn-script-location -e .
-_err_exit $? _err_msg
-
-echo -e "\n***** Installed InvokeAI *****\n"
-
-cp binary_installer/invoke.sh.in ./invoke.sh
-chmod a+rx ./invoke.sh
-echo -e "\n***** Installed invoke launcher script ******\n"
-
-# more cleanup
-rm -rf binary_installer/ installer_files/
-
-# preload the models
-.venv/bin/python3 scripts/configure_invokeai.py
-_err_msg="\n----- model download clone failed -----\n"
-_err_exit $? _err_msg
-deactivate
-
-echo -e "\n***** Finished downloading models *****\n"
-
-echo "All done! Run the command"
-echo "  $scriptdir/invoke.sh"
-echo "to start InvokeAI."
-read -p "Press any key to exit..."
-exit

binary_installer/invoke.bat.in

@@ -1,36 +0,0 @@
-@echo off
-
-PUSHD "%~dp0"
-call .venv\Scripts\activate.bat
-
-echo Do you want to generate images using the
-echo 1. command-line
-echo 2. browser-based UI
-echo OR
-echo 3. open the developer console
-set /p choice="Please enter 1, 2 or 3: "
-if /i "%choice%" == "1" (
-    echo Starting the InvokeAI command-line.
-    .venv\Scripts\python scripts\invoke.py %*
-) else if /i "%choice%" == "2" (
-    echo Starting the InvokeAI browser-based UI.
-    .venv\Scripts\python scripts\invoke.py --web %*
-) else if /i "%choice%" == "3" (
-    echo Developer Console
-    echo Python command is:
-    where python
-    echo Python version is:
-    python --version
-    echo *************************
-    echo You are now in the system shell, with the local InvokeAI Python virtual environment activated,
-    echo so that you can troubleshoot this InvokeAI installation as necessary.
-    echo *************************
-    echo *** Type `exit` to quit this shell and deactivate the Python virtual environment ***
-    call cmd /k
-) else (
-    echo Invalid selection
-    pause
-    exit /b
-)
-
-deactivate

binary_installer/invoke.sh.in

@@ -1,46 +0,0 @@
-#!/usr/bin/env sh
-
-set -eu
-
-. .venv/bin/activate
-
-# set required env var for torch on mac MPS
- if [ "$(uname -s)" == "Darwin" ]; then
-     export PYTORCH_ENABLE_MPS_FALLBACK=1
- fi
-
-echo "Do you want to generate images using the"
-echo "1. command-line"
-echo "2. browser-based UI"
-echo "OR"
-echo "3. open the developer console"
-echo "Please enter 1, 2, or 3:"
-read choice
-
-case $choice in
-    1)
-        printf "\nStarting the InvokeAI command-line..\n";
-        .venv/bin/python scripts/invoke.py $*;
-    ;;
-    2)
-        printf "\nStarting the InvokeAI browser-based UI..\n";
-        .venv/bin/python scripts/invoke.py --web $*;
-    ;;
-    3)
-        printf "\nDeveloper Console:\n";
-        printf "Python command is:\n\t";
-        which python;
-        printf "Python version is:\n\t";
-        python --version;
-        echo "*************************"
-        echo "You are now in your user shell ($SHELL) with the local InvokeAI Python virtual environment activated,";
-        echo "so that you can troubleshoot this InvokeAI installation as necessary.";
-        printf "*************************\n"
-        echo "*** Type \`exit\` to quit this shell and deactivate the Python virtual environment *** ";
-        /usr/bin/env "$SHELL";
-    ;;
-    *)
-        echo "Invalid selection";
-        exit
-    ;;
-esac

binary_installer/readme.txt

@@ -1,17 +0,0 @@
-InvokeAI
-
-Project homepage: https://github.com/invoke-ai/InvokeAI
-
-Installation on Windows:
-    NOTE: You might need to enable Windows Long Paths. If you're not sure,
-    then you almost certainly need to. Simply double-click the 'WinLongPathsEnabled.reg'
-    file. Note that you will need to have admin privileges in order to
-    do this.
-
-    Please double-click the 'install.bat' file (while keeping it inside the invokeAI folder).
-
-Installation on Linux and Mac:
-    Please open the terminal, and run './install.sh' (while keeping it inside the invokeAI folder).
-
-After installation, please run the 'invoke.bat' file (on Windows) or 'invoke.sh'
-file (on Linux/Mac) to start InvokeAI.

binary_installer/requirements.in

@@ -1,33 +0,0 @@
---prefer-binary
---extra-index-url https://download.pytorch.org/whl/torch_stable.html
---extra-index-url https://download.pytorch.org/whl/cu116
---trusted-host https://download.pytorch.org
-accelerate~=0.15
-albumentations
-diffusers[torch]~=0.11
-einops
-eventlet
-flask_cors
-flask_socketio
-flaskwebgui==1.0.3
-getpass_asterisk
-imageio-ffmpeg
-pyreadline3
-realesrgan
-send2trash
-streamlit
-taming-transformers-rom1504
-test-tube
-torch-fidelity
-torch==1.12.1 ; platform_system == 'Darwin'
-torch==1.12.0+cu116 ; platform_system == 'Linux' or platform_system == 'Windows'
-torchvision==0.13.1 ; platform_system == 'Darwin'
-torchvision==0.13.0+cu116 ; platform_system == 'Linux' or platform_system == 'Windows'
-transformers
-picklescan
-https://github.com/openai/CLIP/archive/d50d76daa670286dd6cacf3bcd80b5e4823fc8e1.zip
-https://github.com/invoke-ai/clipseg/archive/1f754751c85d7d4255fa681f4491ff5711c1c288.zip
-https://github.com/invoke-ai/GFPGAN/archive/3f5d2397361199bc4a91c08bb7d80f04d7805615.zip ; platform_system=='Windows'
-https://github.com/invoke-ai/GFPGAN/archive/c796277a1cf77954e5fc0b288d7062d162894248.zip ; platform_system=='Linux' or platform_system=='Darwin'
-https://github.com/Birch-san/k-diffusion/archive/363386981fee88620709cf8f6f2eea167bd6cd74.zip
-https://github.com/invoke-ai/PyPatchMatch/archive/129863937a8ab37f6bbcec327c994c0f932abdbc.zip

docs/features/TEXTUAL_INVERSION.md

@@ -154,8 +154,11 @@ training sets will converge with 2000-3000 steps.
 
 This adjusts how many training images are processed simultaneously in
 each step. Higher values will cause the training process to run more
-quickly, but use more memory. The default size will run with GPUs with
-as little as 12 GB.
+quickly, but use more memory. The default size is selected based on
+whether you have the `xformers` memory-efficient attention library
+installed. If `xformers` is available, the batch size will be 8,
+otherwise 3.  These values were chosen to allow training to run with
+GPUs with as little as 12 GB VRAM.
 
 ### Learning rate
 
@@ -172,8 +175,10 @@ learning rate to improve performance.
 
 ### Use xformers acceleration
 
-This will activate XFormers memory-efficient attention. You need to
-have XFormers installed for this to have an effect.
+This will activate XFormers memory-efficient attention, which will
+reduce memory requirements by half or more and allow you to select a
+higher batch size. You need to have XFormers installed for this to
+have an effect.
 
 ### Learning rate scheduler
 
@@ -250,6 +255,49 @@ invokeai-ti \
        --only_save_embeds
 ```
 
+## Using Distributed Training
+
+If you have multiple GPUs on one machine, or a cluster of GPU-enabled
+machines, you can activate distributed training. See the [HuggingFace
+Accelerate pages](https://huggingface.co/docs/accelerate/index) for
+full information, but the basic recipe is:
+
+1. Enter the InvokeAI developer's console command line by selecting
+option [8] from the `invoke.sh`/`invoke.bat` script.
+
+2. Configurate Accelerate using `accelerate config`:
+```sh
+accelerate config
+```
+This will guide you through the configuration process, including
+specifying how many machines you will run training on and the number
+of GPUs pe rmachine.
+
+You only need to do this once.
+
+3. Launch training from the command line using `accelerate launch`. Be sure
+that your current working directory is the InvokeAI root directory (usually
+named `invokeai` in your home directory):
+
+```sh
+accelerate launch .venv/bin/invokeai-ti \
+       --model=stable-diffusion-1.5 \
+       --resolution=512 \
+       --learnable_property=object \
+       --initializer_token='*' \
+       --placeholder_token='<shraddha>' \
+       --train_data_dir=/home/lstein/invokeai/text-inversion-training-data/shraddha \
+       --output_dir=/home/lstein/invokeai/text-inversion-training/shraddha \
+       --scale_lr \
+       --train_batch_size=10 \
+       --gradient_accumulation_steps=4 \
+       --max_train_steps=2000 \
+       --learning_rate=0.0005 \
+       --lr_scheduler=constant \
+       --mixed_precision=fp16 \
+       --only_save_embeds
+```
+
 ## Using Embeddings
 
 After training completes, the resultant embeddings will be saved into your `$INVOKEAI_ROOT/embeddings/<trigger word>/learned_embeds.bin`.

docs/installation/010_INSTALL_AUTOMATED.md

@@ -417,7 +417,7 @@ Then type the following commands:
 
 === "AMD System"
     ```bash
-    pip install torch torchvision --force-reinstall --extra-index-url https://download.pytorch.org/whl/rocm5.2
+    pip install torch torchvision --force-reinstall --extra-index-url https://download.pytorch.org/whl/rocm5.4.2
     ```
 
 ### Corrupted configuration file

docs/installation/030_INSTALL_CUDA_AND_ROCM.md

@@ -110,7 +110,7 @@ recipes are available
 
 When installing torch and torchvision manually with `pip`, remember to provide
 the argument `--extra-index-url
-https://download.pytorch.org/whl/rocm5.2` as described in the [Manual
+https://download.pytorch.org/whl/rocm5.4.2` as described in the [Manual
 Installation Guide](020_INSTALL_MANUAL.md).
 
 This will be done automatically for you if you use the installer

docs/installation/050_INSTALLING_MODELS.md

@@ -211,6 +211,26 @@ description for the model, whether to make this the default model that
 is loaded at InvokeAI startup time, and whether to replace its
 VAE. Generally the answer to the latter question is "no".
 
+### Specifying a configuration file for legacy checkpoints
+
+Some checkpoint files come with instructions to use a specific .yaml
+configuration file. For InvokeAI load this file correctly, please put
+the config file in the same directory as the corresponding `.ckpt` or
+`.safetensors` file and make sure the file has the same basename as
+the weights file. Here is an example:
+
+```bash
+wonderful-model-v2.ckpt
+wonderful-model-v2.yaml
+```
+
+Similarly, to use a custom VAE, name the VAE like this:
+
+```bash
+wonderful-model-v2.vae.pt
+```
+
+
 ### Converting legacy models into `diffusers`
 
 The CLI `!convert_model` will convert a `.safetensors` or `.ckpt`

installer/lib/installer.py

@@ -241,14 +241,18 @@ class InvokeAiInstance:
 
         from plumbum import FG, local
 
+        # Note that we're installing pinned versions of torch and
+        # torchvision here, which may not correspond to what is
+        # in pyproject.toml. This is a hack to prevent torch 2.0 from
+        # being installed and immediately uninstalled and replaced with 1.13
         pip = local[self.pip]
 
         (
             pip[
                 "install",
                 "--require-virtualenv",
-                "torch",
-                "torchvision",
+                "torch~=1.13.1",
+                "torchvision>=0.14.1",
                 "--force-reinstall",
                 "--find-links" if find_links is not None else None,
                 find_links,
@@ -379,6 +383,9 @@ class InvokeAiInstance:
             shutil.copy(src, dest)
             os.chmod(dest, 0o0755)
 
+        if OS == "Linux":
+            shutil.copy(Path(__file__).parent / '..' / "templates" / "dialogrc", self.runtime / '.dialogrc')
+
     def update(self):
         pass
 

installer/templates/dialogrc

@@ -0,0 +1,27 @@
+# Screen 
+use_shadow = OFF
+use_colors = ON
+screen_color = (BLACK, BLACK, ON)
+
+# Box 
+dialog_color = (YELLOW, BLACK , ON) 
+title_color = (YELLOW, BLACK, ON) 
+border_color = (YELLOW, BLACK, OFF) 
+border2_color = (YELLOW, BLACK, OFF) 
+
+# Button 
+button_active_color = (RED, BLACK, OFF) 
+button_inactive_color = (YELLOW, BLACK, OFF) 
+button_label_active_color = (YELLOW,BLACK,ON) 
+button_label_inactive_color = (YELLOW,BLACK,ON)
+ 
+# Menu box 
+menubox_color = (BLACK, BLACK, ON) 
+menubox_border_color = (YELLOW, BLACK, OFF) 
+menubox_border2_color = (YELLOW, BLACK, OFF) 
+
+# Menu window 
+item_color = (YELLOW, BLACK, OFF) 
+item_selected_color = (BLACK, YELLOW, OFF) 
+tag_key_color = (YELLOW, BLACK, OFF) 
+tag_key_selected_color = (BLACK, YELLOW, OFF)

installer/templates/invoke.sh.in

@@ -1,5 +1,10 @@
 #!/bin/bash
 
+# MIT License
+
+# Coauthored by Lincoln Stein, Eugene Brodsky and Joshua Kimsey
+# Copyright 2023, The InvokeAI Development Team
+
 ####
 # This launch script assumes that:
 # 1. it is located in the runtime directory,
@@ -11,85 +16,168 @@
 
 set -eu
 
-# ensure we're in the correct folder in case user's CWD is somewhere else
+# Ensure we're in the correct folder in case user's CWD is somewhere else
 scriptdir=$(dirname "$0")
 cd "$scriptdir"
 
 . .venv/bin/activate
 
 export INVOKEAI_ROOT="$scriptdir"
+PARAMS=$@
 
-# set required env var for torch on mac MPS
+# Check to see if dialog is installed (it seems to be fairly standard, but good to check regardless) and if the user has passed the --no-tui argument to disable the dialog TUI
+tui=true
+if command -v dialog &>/dev/null; then
+    # This must use $@ to properly loop through the arguments passed by the user
+    for arg in "$@"; do
+        if [ "$arg" == "--no-tui" ]; then
+            tui=false
+            # Remove the --no-tui argument to avoid errors later on when passing arguments to InvokeAI
+            PARAMS=$(echo "$PARAMS" | sed 's/--no-tui//')
+            break
+        fi
+    done
+else
+    tui=false
+fi
+
+# Set required env var for torch on mac MPS
 if [ "$(uname -s)" == "Darwin" ]; then
     export PYTORCH_ENABLE_MPS_FALLBACK=1
 fi
 
-while true
-do
-if [ "$0" != "bash" ]; then
-    echo "Do you want to generate images using the"
-    echo "1. command-line interface"
-    echo "2. browser-based UI"
-    echo "3. run textual inversion training"
-    echo "4. merge models (diffusers type only)"
-    echo "5. download and install models"
-    echo "6. change InvokeAI startup options"
-    echo "7. re-run the configure script to fix a broken install"
-    echo "8. open the developer console"
-    echo "9. update InvokeAI"
-    echo "10. command-line help"
-    echo "Q - Quit"
-    echo ""
-    read -p "Please enter 1-10, Q: [2] " yn
-    choice=${yn:='2'}
-    case $choice in
-        1)
-            echo "Starting the InvokeAI command-line..."
-            invokeai $@
-            ;;
-        2)
-            echo "Starting the InvokeAI browser-based UI..."
-            invokeai --web $@
-            ;;
-        3)
-            echo "Starting Textual Inversion:"
-            invokeai-ti --gui $@
-            ;;
-        4)
-            echo "Merging Models:"
-            invokeai-merge --gui $@
-            ;;
-        5)
-            invokeai-model-install --root ${INVOKEAI_ROOT}
-            ;;
-        6)
-            invokeai-configure --root ${INVOKEAI_ROOT} --skip-sd-weights --skip-support-models
-            ;;
-        7)
-            invokeai-configure --root ${INVOKEAI_ROOT} --yes --default_only
-	    ;;
-	8)
-	    echo "Developer Console:"
-            file_name=$(basename "${BASH_SOURCE[0]}")
-            bash --init-file "$file_name"
-            ;;
-        9)
-	    echo "Update:"
-            invokeai-update
-            ;;
-        10)
-            invokeai --help
-            ;;
-	[qQ])
-            exit 0
-            ;;
-        *)
-            echo "Invalid selection"
-            exit;;
+# Primary function for the case statement to determine user input
+do_choice() {
+    case $1 in
+    1)
+        clear
+        printf "Generate images with a browser-based interface\n"
+        invokeai --web $PARAMS
+        ;;
+    2)
+        clear
+        printf "Generate images using a command-line interface\n"
+        invokeai $PARAMS
+        ;;
+    3)
+        clear
+        printf "Textual inversion training\n"
+        invokeai-ti --gui $PARAMS
+        ;;
+    4)
+        clear
+        printf "Merge models (diffusers type only)\n"
+        invokeai-merge --gui $PARAMS
+        ;;
+    5)
+        clear
+        printf "Download and install models\n"
+        invokeai-model-install --root ${INVOKEAI_ROOT}
+        ;;
+    6)
+        clear
+        printf "Change InvokeAI startup options\n"
+        invokeai-configure --root ${INVOKEAI_ROOT} --skip-sd-weights --skip-support-models
+        ;;
+    7)
+        clear
+        printf "Re-run the configure script to fix a broken install\n"
+        invokeai-configure --root ${INVOKEAI_ROOT} --yes --default_only
+        ;;
+    8)
+        clear
+        printf "Open the developer console\n"
+        file_name=$(basename "${BASH_SOURCE[0]}")
+        bash --init-file "$file_name"
+        ;;
+    9)
+        clear
+        printf "Update InvokeAI\n"
+        invokeai-update
+        ;;
+    10)
+        clear
+        printf "Command-line help\n"
+        invokeai --help
+        ;;
+    "HELP 1")
+        clear
+        printf "Command-line help\n"
+        invokeai --help
+        ;;
+    *)
+        clear
+        printf "Exiting...\n"
+        exit
+        ;;
     esac
+    clear
+}
+
+# Dialog-based TUI for launcing Invoke functions
+do_dialog() {
+    options=(
+        1 "Generate images with a browser-based interface"
+        2 "Generate images using a command-line interface"
+        3 "Textual inversion training"
+        4 "Merge models (diffusers type only)"
+        5 "Download and install models"
+        6 "Change InvokeAI startup options"
+        7 "Re-run the configure script to fix a broken install"
+        8 "Open the developer console"
+        9 "Update InvokeAI")
+
+    choice=$(dialog --clear \
+        --backtitle "\Zb\Zu\Z3InvokeAI" \
+        --colors \
+        --title "What would you like to run?" \
+        --ok-label "Run" \
+        --cancel-label "Exit" \
+        --help-button \
+        --help-label "CLI Help" \
+        --menu "Select an option:" \
+        0 0 0 \
+        "${options[@]}" \
+        2>&1 >/dev/tty) || clear
+    do_choice "$choice"
+    clear
+}
+
+# Command-line interface for launching Invoke functions
+do_line_input() {
+    clear
+    printf " ** For a more attractive experience, please install the 'dialog' utility using your package manager. **\n\n"
+    printf "Do you want to generate images using the\n"
+    printf "1: Browser-based UI\n"
+    printf "2: Command-line interface\n"
+    printf "3: Run textual inversion training\n"
+    printf "4: Merge models (diffusers type only)\n"
+    printf "5: Download and install models\n"
+    printf "6: Change InvokeAI startup options\n"
+    printf "7: Re-run the configure script to fix a broken install\n"
+    printf "8: Open the developer console\n"
+    printf "9: Update InvokeAI\n"
+    printf "10: Command-line help\n"
+    printf "Q: Quit\n\n"
+    read -p "Please enter 1-10, Q: [1] " yn
+    choice=${yn:='1'}
+    do_choice $choice
+    clear
+}
+
+# Main IF statement for launching Invoke with either the TUI or CLI, and for checking if the user is in the developer console
+if [ "$0" != "bash" ]; then
+    while true; do
+        if $tui; then
+            # .dialogrc must be located in the same directory as the invoke.sh script
+            export DIALOGRC="./.dialogrc"
+            do_dialog
+        else
+            do_line_input
+        fi
+    done
 else # in developer console
     python --version
-    echo "Press ^D to exit"
+    printf "Press ^D to exit\n"
     export PS1="(InvokeAI) \u@\h \w> "
 fi
-done

ldm/invoke/CLI.py

@@ -4,6 +4,7 @@ import shlex
 import sys
 import traceback
 from argparse import Namespace
+from packaging import version
 from pathlib import Path
 from typing import Union
 
@@ -1288,8 +1289,27 @@ def check_internet() -> bool:
 
 # This routine performs any patch-ups needed after installation
 def run_patches():
-    # install ckpt configuration files that may have been added to the
-    # distro after original root directory configuration
+    install_missing_config_files()
+    version_file = Path(Globals.root,'.version')
+    if version_file.exists():
+        with open(version_file,'r') as f:
+            root_version = version.parse(f.readline() or 'v2.3.2')
+    else:
+        root_version = version.parse('v2.3.2')
+    app_version = version.parse(ldm.invoke.__version__)
+    if root_version < app_version:
+        try:
+            do_version_update(root_version, ldm.invoke.__version__)
+            with open(version_file,'w') as f:
+                f.write(ldm.invoke.__version__)
+        except:
+            print("** Update failed. Will try again on next launch")
+
+def install_missing_config_files():
+    """
+    install ckpt configuration files that may have been added to the
+    distro after original root directory configuration
+    """
     import invokeai.configs as conf
     from shutil import copyfile
     
@@ -1300,6 +1320,27 @@ def run_patches():
         if not dest.exists():
             copyfile(src,dest)
     
+def do_version_update(root_version: version.Version, app_version: Union[str, version.Version]):
+    """
+    Make any updates to the launcher .sh and .bat scripts that may be needed
+    from release to release. This is not an elegant solution. Instead, the 
+    launcher should be moved into the source tree and installed using pip.
+    """
+    if root_version < version.Version('v2.3.3'):
+        if sys.platform == "linux":
+            print('>> Downloading new version of launcher script and its config file')
+            from ldm.util import download_with_progress_bar
+            url_base = f'https://raw.githubusercontent.com/invoke-ai/InvokeAI/release/v{str(app_version)}/installer/templates/'
+
+            dest = Path(Globals.root,'invoke.sh.in')
+            assert download_with_progress_bar(url_base+'invoke.sh.in',dest)
+            dest.replace(Path(Globals.root,'invoke.sh'))
+            os.chmod(Path(Globals.root,'invoke.sh'), 0o0755)
+
+            dest = Path(Globals.root,'dialogrc')
+            assert download_with_progress_bar(url_base+'dialogrc',dest)
+            dest.replace(Path(Globals.root,'.dialogrc'))
+
 if __name__ == '__main__':
     main()
     

ldm/invoke/_version.py

@@ -1,2 +1,2 @@
 
-__version__='2.3.2'
+__version__='2.3.3-rc2'

ldm/invoke/config/invokeai_update.py

@@ -72,7 +72,7 @@ def main():
         tag = Prompt.ask('Enter an InvokeAI tag or branch name')
 
     print(f':crossed_fingers: Upgrading to [yellow]{tag}[/yellow]')
-    cmd = f'pip install {INVOKE_AI_SRC}/{tag}.zip --use-pep517'
+    cmd = f'pip install {INVOKE_AI_SRC}/{tag}.zip --use-pep517 --upgrade'
     print('')
     print('')
     if os.system(cmd)==0:

ldm/invoke/dynamic_prompts.py

@@ -157,7 +157,7 @@ def _run_invoke(
 ):
     pid = os.getpid()
     logdir.mkdir(parents=True, exist_ok=True)
-    logfile = Path(logdir, f'{time.strftime("%Y-%m-%d-%H:%M:%S")}-pid={pid}.txt')
+    logfile = Path(logdir, f'{time.strftime("%Y-%m-%d_%H-%M-%S")}-pid={pid}.txt')
     print(
         f">> Process {pid} running on GPU {gpu}; logging to {logfile}", file=sys.stderr
     )

ldm/invoke/model_manager.py

@@ -19,7 +19,7 @@ import warnings
 from enum import Enum
 from pathlib import Path
 from shutil import move, rmtree
-from typing import Any, Optional, Union, Callable
+from typing import Any, Callable, Optional, Union
 
 import safetensors
 import safetensors.torch
@@ -35,12 +35,7 @@ from picklescan.scanner import scan_file_path
 from ldm.invoke.devices import CPU_DEVICE
 from ldm.invoke.generator.diffusers_pipeline import StableDiffusionGeneratorPipeline
 from ldm.invoke.globals import Globals, global_cache_dir
-from ldm.util import (
-    ask_user,
-    download_with_resume,
-    instantiate_from_config,
-    url_attachment_name,
-)
+from ldm.util import ask_user, download_with_resume, instantiate_from_config, url_attachment_name
 
 
 class SDLegacyType(Enum):
@@ -282,13 +277,13 @@ class ModelManager(object):
             self.stack.remove(model_name)
         if delete_files:
             if weights:
-                print(f"** deleting file {weights}")
+                print(f"** Deleting file {weights}")
                 Path(weights).unlink(missing_ok=True)
             elif path:
-                print(f"** deleting directory {path}")
+                print(f"** Deleting directory {path}")
                 rmtree(path, ignore_errors=True)
             elif repo_id:
-                print(f"** deleting the cached model directory for {repo_id}")
+                print(f"** Deleting the cached model directory for {repo_id}")
                 self._delete_model_from_cache(repo_id)
 
     def add_model(
@@ -384,15 +379,16 @@ class ModelManager(object):
         if not os.path.isabs(weights):
             weights = os.path.normpath(os.path.join(Globals.root, weights))
 
+        # check whether this is a v2 file and force conversion
+        convert = Globals.ckpt_convert or self.is_v2_config(config)
+
         # if converting automatically to diffusers, then we do the conversion and return
         # a diffusers pipeline
-        if Globals.ckpt_convert:
+        if convert:
             print(
                 f">> Converting legacy checkpoint {model_name} into a diffusers model..."
             )
-            from ldm.invoke.ckpt_to_diffuser import (
-                load_pipeline_from_original_stable_diffusion_ckpt,
-            )
+            from ldm.invoke.ckpt_to_diffuser import load_pipeline_from_original_stable_diffusion_ckpt
 
             self.offload_model(self.current_model)
             if vae_config := self._choose_diffusers_vae(model_name):
@@ -420,11 +416,6 @@ class ModelManager(object):
                 "NOHASH",
             )
 
-        # scan model
-        self.scan_model(model_name, weights)
-
-        print(f">> Loading {model_name} from {weights}")
-
         # for usage statistics
         if self._has_cuda():
             torch.cuda.reset_peak_memory_stats()
@@ -438,10 +429,13 @@ class ModelManager(object):
             weight_bytes = f.read()
         model_hash = self._cached_sha256(weights, weight_bytes)
         sd = None
-        if weights.endswith(".safetensors"):
-            sd = safetensors.torch.load(weight_bytes)
-        else:
+
+        if weights.endswith(".ckpt"):
+            self.scan_model(model_name, weights)
             sd = torch.load(io.BytesIO(weight_bytes), map_location="cpu")
+        else:
+            sd = safetensors.torch.load(weight_bytes)
+
         del weight_bytes
         # merged models from auto11 merge board are flat for some reason
         if "state_dict" in sd:
@@ -464,18 +458,12 @@ class ModelManager(object):
                 vae = os.path.normpath(os.path.join(Globals.root, vae))
             if os.path.exists(vae):
                 print(f"   | Loading VAE weights from: {vae}")
-                vae_ckpt = None
-                vae_dict = None
-                if vae.endswith(".safetensors"):
-                    vae_ckpt = safetensors.torch.load_file(vae)
-                    vae_dict = {k: v for k, v in vae_ckpt.items() if k[0:4] != "loss"}
-                else:
+                if vae.endswith((".ckpt", ".pt")):
+                    self.scan_model(vae, vae)
                     vae_ckpt = torch.load(vae, map_location="cpu")
-                    vae_dict = {
-                        k: v
-                        for k, v in vae_ckpt["state_dict"].items()
-                        if k[0:4] != "loss"
-                    }
+                else:
+                    vae_ckpt = safetensors.torch.load_file(vae)
+                vae_dict = {k: v for k, v in vae_ckpt.items() if k[0:4] != "loss"}
                 model.first_stage_model.load_state_dict(vae_dict, strict=False)
             else:
                 print(f"   | VAE file {vae} not found. Skipping.")
@@ -497,9 +485,9 @@ class ModelManager(object):
 
         print(f">> Loading diffusers model from {name_or_path}")
         if using_fp16:
-            print("  | Using faster float16 precision")
+            print("   | Using faster float16 precision")
         else:
-            print("  | Using more accurate float32 precision")
+            print("   | Using more accurate float32 precision")
 
         # TODO: scan weights maybe?
         pipeline_args: dict[str, Any] = dict(
@@ -551,10 +539,19 @@ class ModelManager(object):
         width = pipeline.unet.config.sample_size * pipeline.vae_scale_factor
         height = width
 
-        print(f"  | Default image dimensions = {width} x {height}")
+        print(f"   | Default image dimensions = {width} x {height}")
 
         return pipeline, width, height, model_hash
 
+    def is_v2_config(self, config: Path) -> bool:
+        try:
+            mconfig = OmegaConf.load(config)
+            return (
+                mconfig["model"]["params"]["unet_config"]["params"]["context_dim"] > 768
+            )
+        except:
+            return False
+
     def model_name_or_path(self, model_name: Union[str, DictConfig]) -> str | Path:
         if isinstance(model_name, DictConfig) or isinstance(model_name, dict):
             mconfig = model_name
@@ -591,13 +588,14 @@ class ModelManager(object):
         if self._has_cuda():
             torch.cuda.empty_cache()
 
+    @classmethod
     def scan_model(self, model_name, checkpoint):
         """
         Apply picklescanner to the indicated checkpoint and issue a warning
         and option to exit if an infected file is identified.
         """
         # scan model
-        print(f">> Scanning Model: {model_name}")
+        print(f"   | Scanning Model: {model_name}")
         scan_result = scan_file_path(checkpoint)
         if scan_result.infected_files != 0:
             if scan_result.infected_files == 1:
@@ -620,7 +618,7 @@ class ModelManager(object):
                     print("### Exiting InvokeAI")
                     sys.exit()
         else:
-            print(">> Model scanned ok")
+            print("   | Model scanned ok")
 
     def import_diffuser_model(
         self,
@@ -731,7 +729,7 @@ class ModelManager(object):
         SDLegacyType.V2_v   (V2 using 'v_prediction' prediction type)
         SDLegacyType.UNKNOWN
         """
-        global_step = checkpoint.get('global_step')
+        global_step = checkpoint.get("global_step")
         state_dict = checkpoint.get("state_dict") or checkpoint
 
         try:
@@ -758,14 +756,14 @@ class ModelManager(object):
             return SDLegacyType.UNKNOWN
 
     def heuristic_import(
-            self,
-            path_url_or_repo: str,
-            convert: bool = False,
-            model_name: str = None,
-            description: str = None,
-            model_config_file: Path = None,
-            commit_to_conf: Path = None,
-            config_file_callback: Callable[[Path],Path] = None,
+        self,
+        path_url_or_repo: str,
+        convert: bool = False,
+        model_name: str = None,
+        description: str = None,
+        model_config_file: Path = None,
+        commit_to_conf: Path = None,
+        config_file_callback: Callable[[Path], Path] = None,
     ) -> str:
         """
         Accept a string which could be:
@@ -800,19 +798,20 @@ class ModelManager(object):
         print(f">> Probing {thing} for import")
 
         if thing.startswith(("http:", "https:", "ftp:")):
-            print(f"  | {thing} appears to be a URL")
+            print(f"   | {thing} appears to be a URL")
             model_path = self._resolve_path(
                 thing, "models/ldm/stable-diffusion-v1"
             )  # _resolve_path does a download if needed
             is_temporary = True
+
         elif Path(thing).is_file() and thing.endswith((".ckpt", ".safetensors")):
             if Path(thing).stem in ["model", "diffusion_pytorch_model"]:
                 print(
-                    f"  | {Path(thing).name} appears to be part of a diffusers model. Skipping import"
+                    f"   | {Path(thing).name} appears to be part of a diffusers model. Skipping import"
                 )
                 return
             else:
-                print(f"  | {thing} appears to be a checkpoint file on disk")
+                print(f"   | {thing} appears to be a checkpoint file on disk")
                 model_path = self._resolve_path(thing, "models/ldm/stable-diffusion-v1")
 
         elif Path(thing).is_dir() and Path(thing, "model_index.json").exists():
@@ -839,10 +838,10 @@ class ModelManager(object):
                     Path(thing).rglob("*.safetensors")
                 ):
                     if model_name := self.heuristic_import(
-                            str(m),
-                            convert,
-                            commit_to_conf=commit_to_conf,
-                            config_file_callback=config_file_callback,
+                        str(m),
+                        convert,
+                        commit_to_conf=commit_to_conf,
+                        config_file_callback=config_file_callback,
                     ):
                         print(f" >> {model_name} successfully imported")
                 return model_name
@@ -869,57 +868,67 @@ class ModelManager(object):
             return model_path.stem
 
         # another round of heuristics to guess the correct config file.
-        checkpoint = (
-            safetensors.torch.load_file(model_path)
-            if model_path.suffix == ".safetensors"
-            else torch.load(model_path)
-        )
+        checkpoint = None
+        if model_path.suffix.endswith((".ckpt", ".pt")):
+            self.scan_model(model_path, model_path)
+            checkpoint = torch.load(model_path)
+        else:
+            checkpoint = safetensors.torch.load_file(model_path)
         # additional probing needed if no config file provided
         if model_config_file is None:
-            model_type = self.probe_model_type(checkpoint)
-            if model_type == SDLegacyType.V1:
-                print("  | SD-v1 model detected")
-                model_config_file = Path(
-                    Globals.root, "configs/stable-diffusion/v1-inference.yaml"
-                )
-            elif model_type == SDLegacyType.V1_INPAINT:
-                print("  | SD-v1 inpainting model detected")
-                model_config_file = Path(
-                    Globals.root, "configs/stable-diffusion/v1-inpainting-inference.yaml"
-                )
-            elif model_type == SDLegacyType.V2_v:
-                print(
-                    "  | SD-v2-v model detected"
-                )
-                model_config_file = Path(
-                    Globals.root, "configs/stable-diffusion/v2-inference-v.yaml"
-                )
-            elif model_type == SDLegacyType.V2_e:
-                print(
-                    "  | SD-v2-e model detected"
-                )
-                model_config_file = Path(
-                    Globals.root, "configs/stable-diffusion/v2-inference.yaml"
-                )
-            elif model_type == SDLegacyType.V2:
-                print(
-                    f"** {thing} is a V2 checkpoint file, but its parameterization cannot be determined. Please provide configuration file path."
-                )
+            # Is there a like-named .yaml file in the same directory as the
+            # weights file? If so, we treat this as our model
+            if model_path.with_suffix(".yaml").exists():
+                model_config_file = model_path.with_suffix(".yaml")
+                print(f"   | Using config file {model_config_file.name}")
             else:
-                print(
-                    f"** {thing} is a legacy checkpoint file but not a known Stable Diffusion model. Please provide configuration file path."
-                )
+                model_type = self.probe_model_type(checkpoint)
+                if model_type == SDLegacyType.V1:
+                    print("   | SD-v1 model detected")
+                    model_config_file = Path(
+                        Globals.root, "configs/stable-diffusion/v1-inference.yaml"
+                    )
+                elif model_type == SDLegacyType.V1_INPAINT:
+                    print("   | SD-v1 inpainting model detected")
+                    model_config_file = Path(
+                        Globals.root,
+                        "configs/stable-diffusion/v1-inpainting-inference.yaml",
+                    )
+                elif model_type == SDLegacyType.V2_v:
+                    print("   | SD-v2-v model detected")
+                    model_config_file = Path(
+                        Globals.root, "configs/stable-diffusion/v2-inference-v.yaml"
+                    )
+                elif model_type == SDLegacyType.V2_e:
+                    print("   | SD-v2-e model detected")
+                    model_config_file = Path(
+                        Globals.root, "configs/stable-diffusion/v2-inference.yaml"
+                    )
+                elif model_type == SDLegacyType.V2:
+                    print(
+                        f"** {thing} is a V2 checkpoint file, but its parameterization cannot be determined. Please provide configuration file path."
+                    )
+                else:
+                    print(
+                        f"** {thing} is a legacy checkpoint file but not a known Stable Diffusion model. Please provide configuration file path."
+                    )
 
             if not model_config_file and config_file_callback:
                 model_config_file = config_file_callback(model_path)
             if not model_config_file:
                 return
 
-        if model_config_file.name.startswith('v2'):
+        if self.is_v2_config(model_config_file):
             convert = True
-            print(
-                "  | This SD-v2 model will be converted to diffusers format for use"
-            )
+            print("   | This SD-v2 model will be converted to diffusers format for use")
+
+        # look for a custom vae
+        vae_path = None
+        for suffix in ["pt", "ckpt", "safetensors"]:
+            if (model_path.with_suffix(f".vae.{suffix}")).exists():
+                vae_path = model_path.with_suffix(f".vae.{suffix}")
+                print(f"   | Using VAE file {vae_path.name}")
+        vae = None if vae_path else dict(repo_id="stabilityai/sd-vae-ft-mse")
 
         if convert:
             diffuser_path = Path(
@@ -928,11 +937,13 @@ class ModelManager(object):
             model_name = self.convert_and_import(
                 model_path,
                 diffusers_path=diffuser_path,
-                vae=dict(repo_id="stabilityai/sd-vae-ft-mse"),
+                vae=vae,
+                vae_path=vae_path,
                 model_name=model_name,
                 model_description=description,
                 original_config_file=model_config_file,
                 commit_to_conf=commit_to_conf,
+                scan_needed=False,
             )
             # in the event that this file was downloaded automatically prior to conversion
             # we do not keep the original .ckpt/.safetensors around
@@ -945,7 +956,8 @@ class ModelManager(object):
                 model_name=model_name,
                 model_description=description,
                 vae=str(
-                    Path(
+                    vae_path
+                    or Path(
                         Globals.root,
                         "models/ldm/stable-diffusion-v1/vae-ft-mse-840000-ema-pruned.ckpt",
                     )
@@ -962,9 +974,11 @@ class ModelManager(object):
         diffusers_path: Path,
         model_name=None,
         model_description=None,
-        vae=None,
+        vae: dict = None,
+        vae_path: Path = None,
         original_config_file: Path = None,
         commit_to_conf: Path = None,
+        scan_needed: bool = True,
     ) -> str:
         """
         Convert a legacy ckpt weights file to diffuser model and import
@@ -978,7 +992,7 @@ class ModelManager(object):
 
         new_config = None
 
-        from ldm.invoke.ckpt_to_diffuser import convert_ckpt_to_diffuser
+        from ldm.invoke.ckpt_to_diffuser import convert_ckpt_to_diffusers
 
         if diffusers_path.exists():
             print(
@@ -993,17 +1007,19 @@ class ModelManager(object):
             # By passing the specified VAE to the conversion function, the autoencoder
             # will be built into the model rather than tacked on afterward via the config file
             vae_model = self._load_vae(vae) if vae else None
-            convert_ckpt_to_diffuser(
+            convert_ckpt_to_diffusers(
                 ckpt_path,
                 diffusers_path,
                 extract_ema=True,
                 original_config_file=original_config_file,
                 vae=vae_model,
+                vae_path=str(vae_path) if vae_path else None,
+                scan_needed=scan_needed,
             )
             print(
-                f"  | Success. Optimized model is now located at {str(diffusers_path)}"
+                f"   | Success. Optimized model is now located at {str(diffusers_path)}"
             )
-            print(f"  | Writing new config file entry for {model_name}")
+            print(f"   | Writing new config file entry for {model_name}")
             new_config = dict(
                 path=str(diffusers_path),
                 description=model_description,
@@ -1050,7 +1066,7 @@ class ModelManager(object):
         # In the event that the original entry is using a custom ckpt VAE, we try to
         # map that VAE onto a diffuser VAE using a hard-coded dictionary.
         # I would prefer to do this differently: We load the ckpt model into memory, swap the
-        # VAE in memory, and then pass that to convert_ckpt_to_diffuser() so that the swapped
+        # VAE in memory, and then pass that to convert_ckpt_to_diffusers() so that the swapped
         # VAE is built into the model. However, when I tried this I got obscure key errors.
         if vae:
             return vae
@@ -1136,14 +1152,14 @@ class ModelManager(object):
         legacy_locations = [
             Path(
                 models_dir,
-                "CompVis/stable-diffusion-safety-checker/models--CompVis--stable-diffusion-safety-checker"
+                "CompVis/stable-diffusion-safety-checker/models--CompVis--stable-diffusion-safety-checker",
             ),
             Path("bert-base-uncased/models--bert-base-uncased"),
             Path(
                 "openai/clip-vit-large-patch14/models--openai--clip-vit-large-patch14"
             ),
         ]
-        legacy_locations.extend(list(global_cache_dir("diffusers").glob('*')))
+        legacy_locations.extend(list(global_cache_dir("diffusers").glob("*")))
         legacy_layout = False
         for model in legacy_locations:
             legacy_layout = legacy_layout or model.exists()
@@ -1187,7 +1203,7 @@ class ModelManager(object):
                         source.unlink()
                 else:
                     move(source, dest)
-                    
+
         # now clean up by removing any empty directories
         empty = [
             root
@@ -1293,7 +1309,7 @@ class ModelManager(object):
             with open(hashpath) as f:
                 hash = f.read()
             return hash
-        print("  | Calculating sha256 hash of model files")
+        print("   | Calculating sha256 hash of model files")
         tic = time.time()
         sha = hashlib.sha256()
         count = 0
@@ -1305,7 +1321,7 @@ class ModelManager(object):
                         sha.update(chunk)
         hash = sha.hexdigest()
         toc = time.time()
-        print(f"  | sha256 = {hash} ({count} files hashed in", "%4.2fs)" % (toc - tic))
+        print(f"   | sha256 = {hash} ({count} files hashed in", "%4.2fs)" % (toc - tic))
         with open(hashpath, "w") as f:
             f.write(hash)
         return hash
@@ -1350,12 +1366,12 @@ class ModelManager(object):
             local_files_only=not Globals.internet_available,
         )
 
-        print(f"  | Loading diffusers VAE from {name_or_path}")
+        print(f"   | Loading diffusers VAE from {name_or_path}")
         if using_fp16:
             vae_args.update(torch_dtype=torch.float16)
             fp_args_list = [{"revision": "fp16"}, {}]
         else:
-            print("  | Using more accurate float32 precision")
+            print("   | Using more accurate float32 precision")
             fp_args_list = [{}]
 
         vae = None
@@ -1396,7 +1412,7 @@ class ModelManager(object):
                     hashes_to_delete.add(revision.commit_hash)
         strategy = cache_info.delete_revisions(*hashes_to_delete)
         print(
-            f"** deletion of this model is expected to free {strategy.expected_freed_size_str}"
+            f"** Deletion of this model is expected to free {strategy.expected_freed_size_str}"
         )
         strategy.execute()
 

ldm/invoke/pngwriter.py

@@ -30,14 +30,17 @@ class PngWriter:
                 prefix = self._unused_prefix()
             else:
                 with open(next_prefix_file,'r') as file:
-                    prefix=int(file.readline() or int(self._unused_prefix())-1)
-                    prefix+=1
+                    prefix = 0
+                    try:
+                        prefix=int(file.readline())
+                    except (TypeError, ValueError):
+                        prefix=self._unused_prefix()
             with open(next_prefix_file,'w') as file:
-                file.write(str(prefix))
+                file.write(str(prefix+1))
         return f'{prefix:06}'
 
     # gives the next unique prefix in outdir
-    def _unused_prefix(self):
+    def _unused_prefix(self)->int:
         # sort reverse alphabetically until we find max+1
         dirlist = sorted(os.listdir(self.outdir), reverse=True)
         # find the first filename that matches our pattern or return 000000.0.png
@@ -45,8 +48,7 @@ class PngWriter:
             (f for f in dirlist if re.match('^(\d+)\..*\.png', f)),
             '0000000.0.png',
         )
-        basecount = int(existing_name.split('.', 1)[0]) + 1
-        return f'{basecount:06}'
+        return int(existing_name.split('.', 1)[0]) + 1
 
     # saves image named _image_ to outdir/name, writing metadata from prompt
     # returns full path of output

ldm/invoke/training/textual_inversion.py

@@ -17,6 +17,7 @@ from pathlib import Path
 from typing import List, Tuple
 
 import npyscreen
+from diffusers.utils.import_utils import is_xformers_available
 from npyscreen import widget
 from omegaconf import OmegaConf
 
@@ -29,7 +30,7 @@ from ldm.invoke.training.textual_inversion_training import (
 TRAINING_DATA = "text-inversion-training-data"
 TRAINING_DIR = "text-inversion-output"
 CONF_FILE = "preferences.conf"
-
+XFORMERS_AVAILABLE = is_xformers_available()
 
 class textualInversionForm(npyscreen.FormMultiPageAction):
     resolutions = [512, 768, 1024]
@@ -178,7 +179,7 @@ class textualInversionForm(npyscreen.FormMultiPageAction):
             out_of=10000,
             step=500,
             lowest=1,
-            value=saved_args.get("max_train_steps", 3000),
+            value=saved_args.get("max_train_steps", 2500),
             scroll_exit=True,
         )
         self.train_batch_size = self.add_widget_intelligent(
@@ -187,7 +188,7 @@ class textualInversionForm(npyscreen.FormMultiPageAction):
             out_of=50,
             step=1,
             lowest=1,
-            value=saved_args.get("train_batch_size", 8),
+            value=saved_args.get("train_batch_size", 8 if XFORMERS_AVAILABLE else 3),
             scroll_exit=True,
         )
         self.gradient_accumulation_steps = self.add_widget_intelligent(
@@ -225,7 +226,7 @@ class textualInversionForm(npyscreen.FormMultiPageAction):
         self.enable_xformers_memory_efficient_attention = self.add_widget_intelligent(
             npyscreen.Checkbox,
             name="Use xformers acceleration",
-            value=saved_args.get("enable_xformers_memory_efficient_attention", False),
+            value=saved_args.get("enable_xformers_memory_efficient_attention", XFORMERS_AVAILABLE),
             scroll_exit=True,
         )
         self.lr_scheduler = self.add_widget_intelligent(
@@ -428,8 +429,7 @@ def do_front_end(args: Namespace):
             print(str(e))
             print("** DETAILS:")
             print(traceback.format_exc())
-
-
+            
 def main():
     args = parse_args()
     global_set_root(args.root_dir or Globals.root)

ldm/invoke/training/textual_inversion_training.py

@@ -67,7 +67,7 @@ else:
         "nearest": PIL.Image.NEAREST,
     }
 # ------------------------------------------------------------------------------
-
+XFORMERS_AVAILABLE = is_xformers_available
 
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
 check_min_version("0.10.0.dev0")
@@ -227,7 +227,7 @@ def parse_args():
     training_group.add_argument(
         "--train_batch_size",
         type=int,
-        default=16,
+        default=8 if XFORMERS_AVAILABLE else 3,
         help="Batch size (per device) for the training dataloader.",
     )
     training_group.add_argument("--num_train_epochs", type=int, default=100)
@@ -324,6 +324,7 @@ def parse_args():
     parser.add_argument(
         "--enable_xformers_memory_efficient_attention",
         action="store_true",
+        default=XFORMERS_AVAILABLE,
         help="Whether or not to use xformers.",
     )
 
@@ -536,7 +537,7 @@ def do_textual_inversion_training(
     seed: int = None,
     resolution: int = 512,
     center_crop: bool = False,
-    train_batch_size: int = 16,
+    train_batch_size: int = 4,
     num_train_epochs: int = 100,
     max_train_steps: int = 5000,
     gradient_accumulation_steps: int = 1,

ldm/util.py

@@ -329,7 +329,7 @@ def download_with_resume(url: str, dest: Path, access_token: str = None) -> Path
         resp = requests.get(url, headers=header, stream=True) # new request with range
 
     if exist_size > content_length:
-        print('* corrupt existing file found. re-downloading')
+        print(f'* corrupt existing file found (existing_size={exist_size}, content_length={content_length}). re-downloading')
         os.remove(dest)
         exist_size = 0
 
@@ -341,15 +341,12 @@ def download_with_resume(url: str, dest: Path, access_token: str = None) -> Path
     elif resp.status_code == 206 or exist_size > 0:
         print(f"* {dest}: partial file found. Resuming...")
     elif resp.status_code != 200:
-        print(f"** An error occurred during downloading {dest}: {resp.reason}")
+        print(f"** An error occurred while downloading {url}: {resp.reason}")
+        return None
     else:
         print(f"* {dest}: Downloading...")
 
     try:
-        if content_length < 2000:
-            print(f"*** ERROR DOWNLOADING {url}: {resp.text}")
-            return None
-
         with open(dest, open_mode) as file, tqdm(
                 desc=str(dest),
                 initial=exist_size,

pyproject.toml

@@ -58,6 +58,7 @@ dependencies = [
   "pillow",
   "prompt-toolkit",
   "pudb",
+  "packaging",
   "pypatchmatch",
   "pyreadline3",
   "pytorch-lightning==1.7.7",
@@ -70,7 +71,7 @@ dependencies = [
   "taming-transformers-rom1504",
   "test-tube>=0.7.5",
   "torch-fidelity",
-  "torch>=1.13.1",
+  "torch~=1.13.1",
   "torchmetrics",
   "torchvision>=0.14.1",
   "transformers~=4.26",
@@ -147,7 +148,7 @@ version = {attr = "ldm.invoke.__version__"}
 
 [tool.setuptools.package-data]
 "invokeai.assets.web" = ["**.png"]
-"invokeai.configs" = ["**.example", "**.txt", "**.yaml"]
+"invokeai.configs" = ["**.example", "**.txt", "**.yaml", "**/*.yaml"]
 "invokeai.frontend.dist" = ["**"]
 
 [tool.black]