@@ -44,9 +43,18 @@ The installer will take care of whatever is needed. If you face any problems, yo ### User experience - **Hassle-free installation**: Does not require technical knowledge, does not require pre-installed software. Just download and run! - **Clutter-free UI**: A friendly and simple UI, while providing a lot of powerful features. +- **Task Queue**: Queue up all your ideas, without waiting for the current task to finish. +- **Intelligent Model Detection**: Automatically figures out the YAML config file to use for the chosen model (via a models database). +- **Live Preview**: See the image as the AI is drawing it. +- **Image Modifiers**: A library of *modifier tags* like *"Realistic"*, *"Pencil Sketch"*, *"ArtStation"* etc. Experiment with various styles quickly. +- **Multiple Prompts File**: Queue multiple prompts by entering one prompt per line, or by running a text file. +- **Save generated images to disk**: Save your images to your PC! +- **UI Themes**: Customize the program to your liking. +- **Organize your models into sub-folders** ### Image generation - **Supports**: "*Text to Image*" and "*Image to Image*". +- **14 Samplers**: `ddim`, `plms`, `heun`, `euler`, `euler_a`, `dpm2`, `dpm2_a`, `lms`, `dpm_solver_stability`, `dpmpp_2s_a`, `dpmpp_2m`, `dpmpp_sde`, `dpm_fast`, `dpm_adaptive` - **In-Painting**: Specify areas of your image to paint into. - **Simple Drawing Tool**: Draw basic images to guide the AI, without needing an external drawing program. - **Face Correction (GFPGAN)** @@ -56,7 +64,6 @@ The installer will take care of whatever is needed. If you face any problems, yo - **Attention/Emphasis**: () in the prompt increases the model's attention to enclosed words, and [] decreases it. - **Weighted Prompts**: Use weights for specific words in your prompt to change their importance, e.g. `red:2.4 dragon:1.2`. - **Prompt Matrix**: Quickly create multiple variations of your prompt, e.g. `a photograph of an astronaut riding a horse | illustration | cinematic lighting`. -- **Lots of Samplers**: ddim, plms, heun, euler, euler_a, dpm2, dpm2_a, lms. - **1-click Upscale/Face Correction**: Upscale or correct an image after it has been generated. - **Make Similar Images**: Click to generate multiple variations of a generated image. - **NSFW Setting**: A setting in the UI to control *NSFW content*. @@ -64,13 +71,15 @@ The installer will take care of whatever is needed. If you face any problems, yo ### Advanced features - **Custom Models**: Use your own `.ckpt` or `.safetensors` file, by placing it inside the `models/stable-diffusion` folder! -- **Stable Diffusion 2.0 support (experimental)**: available in beta channel. +- **Stable Diffusion 2.1 support** +- **Merge Models** - **Use custom VAE models** - **Use pre-trained Hypernetworks** - **UI Plugins**: Choose from a growing list of [community-generated UI plugins](https://github.com/cmdr2/stable-diffusion-ui/wiki/UI-Plugins), or write your own plugin to add features to the project! ### Performance and security -- **Low Memory Usage**: Creates 512x512 images with less than 4GB of GPU RAM! +- **Fast**: Creates a 512x512 image with euler_a in 5 seconds, on an NVIDIA 3060 12GB. +- **Low Memory Usage**: Create 512x512 images with less than 3 GB of GPU RAM, and 768x768 images with less than 4 GB of GPU RAM! - **Use CPU setting**: If you don't have a compatible graphics card, but still want to run it on your CPU. - **Multi-GPU support**: Automatically spreads your tasks across multiple GPUs (if available), for faster performance! - **Auto scan for malicious models**: Uses picklescan to prevent malicious models. @@ -78,23 +87,17 @@ The installer will take care of whatever is needed. If you face any problems, yo - **Auto-updater**: Gets you the latest improvements and bug-fixes to a rapidly evolving project. - **Developer Console**: A developer-mode for those who want to modify their Stable Diffusion code, and edit the conda environment. -### Usability: -- **Live Preview**: See the image as the AI is drawing it. -- **Task Queue**: Queue up all your ideas, without waiting for the current task to finish. -- **Image Modifiers**: A library of *modifier tags* like *"Realistic"*, *"Pencil Sketch"*, *"ArtStation"* etc. Experiment with various styles quickly. -- **Multiple Prompts File**: Queue multiple prompts by entering one prompt per line, or by running a text file. -- **Save generated images to disk**: Save your images to your PC! -- **UI Themes**: Customize the program to your liking. - **(and a lot more)** ---- ## Easy for new users: - + + ## Powerful features for advanced users: - + + ## Live Preview Useful for judging (and stopping) an image quickly, without waiting for it to finish rendering. @@ -102,7 +105,9 @@ Useful for judging (and stopping) an image quickly, without waiting for it to fi  ## Task Queue - + + + # System Requirements 1. Windows 10/11, or Linux. Experimental support for Mac is coming soon. diff --git a/scripts/Developer Console.cmd b/scripts/Developer Console.cmd index 750e4311..921a9dca 100644 --- a/scripts/Developer Console.cmd +++ b/scripts/Developer Console.cmd @@ -23,23 +23,20 @@ call conda --version echo. -@rem activate the environment -call conda activate .\stable-diffusion\env +@rem activate the legacy environment (if present) and set PYTHONPATH +if exist "installer_files\env" ( + set PYTHONPATH=%cd%\installer_files\env\lib\site-packages +) +if exist "stable-diffusion\env" ( + call conda activate .\stable-diffusion\env + set PYTHONPATH=%cd%\stable-diffusion\env\lib\site-packages +) call where python call python --version -@rem set the PYTHONPATH -cd stable-diffusion -set SD_DIR=%cd% - -cd env\lib\site-packages -set PYTHONPATH=%SD_DIR%;%cd% -cd ..\..\.. echo PYTHONPATH=%PYTHONPATH% -cd .. - @rem done echo. diff --git a/scripts/bootstrap.bat b/scripts/bootstrap.bat index cb0909d0..c58ddcce 100644 --- a/scripts/bootstrap.bat +++ b/scripts/bootstrap.bat @@ -24,7 +24,7 @@ if exist "%INSTALL_ENV_DIR%" set PATH=%INSTALL_ENV_DIR%;%INSTALL_ENV_DIR%\Librar set PACKAGES_TO_INSTALL= if not exist "%LEGACY_INSTALL_ENV_DIR%\etc\profile.d\conda.sh" ( - if not exist "%INSTALL_ENV_DIR%\etc\profile.d\conda.sh" set PACKAGES_TO_INSTALL=%PACKAGES_TO_INSTALL% conda + if not exist "%INSTALL_ENV_DIR%\etc\profile.d\conda.sh" set PACKAGES_TO_INSTALL=%PACKAGES_TO_INSTALL% conda python=3.8.5 ) call git --version >.tmp1 2>.tmp2 diff --git a/scripts/bootstrap.sh b/scripts/bootstrap.sh index 6a929133..6aba4480 100755 --- a/scripts/bootstrap.sh +++ b/scripts/bootstrap.sh @@ -21,6 +21,7 @@ OS_ARCH=$(uname -m) case "${OS_ARCH}" in x86_64*) OS_ARCH="64";; arm64*) OS_ARCH="arm64";; + aarch64*) OS_ARCH="arm64";; *) echo "Unknown system architecture: $OS_ARCH! This script runs only on x86_64 or arm64" && exit esac @@ -39,7 +40,7 @@ if [ -e "$INSTALL_ENV_DIR" ]; then export PATH="$INSTALL_ENV_DIR/bin:$PATH"; fi PACKAGES_TO_INSTALL="" -if [ ! -e "$LEGACY_INSTALL_ENV_DIR/etc/profile.d/conda.sh" ] && [ ! -e "$INSTALL_ENV_DIR/etc/profile.d/conda.sh" ]; then PACKAGES_TO_INSTALL="$PACKAGES_TO_INSTALL conda"; fi +if [ ! -e "$LEGACY_INSTALL_ENV_DIR/etc/profile.d/conda.sh" ] && [ ! -e "$INSTALL_ENV_DIR/etc/profile.d/conda.sh" ]; then PACKAGES_TO_INSTALL="$PACKAGES_TO_INSTALL conda python=3.8.5"; fi if ! hash "git" &>/dev/null; then PACKAGES_TO_INSTALL="$PACKAGES_TO_INSTALL git"; fi if "$MAMBA_ROOT_PREFIX/micromamba" --version &>/dev/null; then umamba_exists="T"; fi diff --git a/scripts/check_modules.py b/scripts/check_modules.py new file mode 100644 index 00000000..416ad851 --- /dev/null +++ b/scripts/check_modules.py @@ -0,0 +1,13 @@ +''' +This script checks if the given modules exist +''' + +import sys +import pkgutil + +modules = sys.argv[1:] +missing_modules = [] +for m in modules: + if pkgutil.find_loader(m) is None: + print('module', m, 'not found') + exit(1) diff --git a/scripts/developer_console.sh b/scripts/developer_console.sh index 49e71b34..73972568 100755 --- a/scripts/developer_console.sh +++ b/scripts/developer_console.sh @@ -26,21 +26,23 @@ if [ "$0" == "bash" ]; then echo "" - # activate the environment - CONDA_BASEPATH=$(conda info --base) - source "$CONDA_BASEPATH/etc/profile.d/conda.sh" # otherwise conda complains about 'shell not initialized' (needed when running in a script) + # activate the legacy environment (if present) and set PYTHONPATH + if [ -e "installer_files/env" ]; then + export PYTHONPATH="$(pwd)/installer_files/env/lib/python3.8/site-packages" + fi + if [ -e "stable-diffusion/env" ]; then + CONDA_BASEPATH=$(conda info --base) + source "$CONDA_BASEPATH/etc/profile.d/conda.sh" # otherwise conda complains about 'shell not initialized' (needed when running in a script) - conda activate ./stable-diffusion/env + conda activate ./stable-diffusion/env + + export PYTHONPATH="$(pwd)/stable-diffusion/env/lib/python3.8/site-packages" + fi which python python --version - # set the PYTHONPATH - cd stable-diffusion - SD_PATH=`pwd` - export PYTHONPATH="$SD_PATH:$SD_PATH/env/lib/python3.8/site-packages" echo "PYTHONPATH=$PYTHONPATH" - cd .. # done diff --git a/scripts/on_env_start.bat b/scripts/on_env_start.bat index b18b4f4e..13d54161 100644 --- a/scripts/on_env_start.bat +++ b/scripts/on_env_start.bat @@ -53,6 +53,7 @@ if "%update_branch%"=="" ( @xcopy sd-ui-files\ui ui /s /i /Y /q @copy sd-ui-files\scripts\on_sd_start.bat scripts\ /Y @copy sd-ui-files\scripts\bootstrap.bat scripts\ /Y +@copy sd-ui-files\scripts\check_modules.py scripts\ /Y @copy "sd-ui-files\scripts\Start Stable Diffusion UI.cmd" . /Y @copy "sd-ui-files\scripts\Developer Console.cmd" . /Y diff --git a/scripts/on_env_start.sh b/scripts/on_env_start.sh index 185e4fd4..cac307c2 100755 --- a/scripts/on_env_start.sh +++ b/scripts/on_env_start.sh @@ -37,6 +37,7 @@ rm -rf ui cp -Rf sd-ui-files/ui . cp sd-ui-files/scripts/on_sd_start.sh scripts/ cp sd-ui-files/scripts/bootstrap.sh scripts/ +cp sd-ui-files/scripts/check_modules.py scripts/ cp sd-ui-files/scripts/start.sh . cp sd-ui-files/scripts/developer_console.sh . diff --git a/scripts/on_sd_start.bat b/scripts/on_sd_start.bat index f4afb650..5b778aa3 100644 --- a/scripts/on_sd_start.bat +++ b/scripts/on_sd_start.bat @@ -5,11 +5,20 @@ @copy sd-ui-files\scripts\on_env_start.bat scripts\ /Y @copy sd-ui-files\scripts\bootstrap.bat scripts\ /Y +@copy sd-ui-files\scripts\check_modules.py scripts\ /Y if exist "%cd%\profile" ( set USERPROFILE=%cd%\profile ) +@rem set the correct installer path (current vs legacy) +if exist "%cd%\installer_files\env" ( + set INSTALL_ENV_DIR=%cd%\installer_files\env +) +if exist "%cd%\stable-diffusion\env" ( + set INSTALL_ENV_DIR=%cd%\stable-diffusion\env +) + @mkdir tmp @set TMP=%cd%\tmp @set TEMP=%cd%\tmp @@ -27,150 +36,117 @@ if exist "Open Developer Console.cmd" del "Open Developer Console.cmd" @call python -c "import os; import shutil; frm = 'sd-ui-files\\ui\\hotfix\\9c24e6cd9f499d02c4f21a033736dabd365962dc80fe3aeb57a8f85ea45a20a3.26fead7ea4f0f843f6eb4055dfd25693f1a71f3c6871b184042d4b126244e142'; dst = os.path.join(os.path.expanduser('~'), '.cache', 'huggingface', 'transformers', '9c24e6cd9f499d02c4f21a033736dabd365962dc80fe3aeb57a8f85ea45a20a3.26fead7ea4f0f843f6eb4055dfd25693f1a71f3c6871b184042d4b126244e142'); shutil.copyfile(frm, dst) if os.path.exists(dst) else print(''); print('Hotfixed broken JSON file from OpenAI');" -if NOT DEFINED test_sd2 set test_sd2=N +@rem create the stable-diffusion folder, to work with legacy installations +if not exist "stable-diffusion" mkdir stable-diffusion +cd stable-diffusion -@>nul findstr /m "sd_git_cloned" scripts\install_status.txt -@if "%ERRORLEVEL%" EQU "0" ( - @echo "Stable Diffusion's git repository was already installed. Updating.." - - @cd stable-diffusion - - @call git remote set-url origin https://github.com/easydiffusion/diffusion-kit.git - - @call git reset --hard - @call git pull - - if "%test_sd2%" == "N" ( - @call git -c advice.detachedHead=false checkout 7f32368ed1030a6e710537047bacd908adea183a - ) - if "%test_sd2%" == "Y" ( - @call git -c advice.detachedHead=false checkout 733a1f6f9cae9b9a9b83294bf3281b123378cb1f - ) - - @cd .. -) else ( - @echo. & echo "Downloading Stable Diffusion.." & echo. - - @call git clone https://github.com/easydiffusion/diffusion-kit.git stable-diffusion && ( - @echo sd_git_cloned >> scripts\install_status.txt - ) || ( - @echo "Error downloading Stable Diffusion. Sorry about that, please try to:" & echo " 1. Run this installer again." & echo " 2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting" & echo " 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" & echo " 4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues" & echo "Thanks!" - pause - @exit /b - ) - - @cd stable-diffusion - @call git -c advice.detachedHead=false checkout 7f32368ed1030a6e710537047bacd908adea183a - - @cd .. +@rem activate the old stable-diffusion env, if it exists +if exist "env" ( + call conda activate .\env ) -@cd stable-diffusion +@rem disable the legacy src and ldm folder (otherwise this prevents installing gfpgan and realesrgan) +if exist src rename src src-old +if exist ldm rename ldm ldm-old -@>nul findstr /m "conda_sd_env_created" ..\scripts\install_status.txt -@if "%ERRORLEVEL%" EQU "0" ( - @echo "Packages necessary for Stable Diffusion were already installed" +if not exist "..\models\stable-diffusion" mkdir "..\models\stable-diffusion" +if not exist "..\models\gfpgan" mkdir "..\models\gfpgan" +if not exist "..\models\realesrgan" mkdir "..\models\realesrgan" +if not exist "..\models\vae" mkdir "..\models\vae" - @call conda activate .\env +@rem migrate the legacy models to the correct path (if already downloaded) +if exist "sd-v1-4.ckpt" move sd-v1-4.ckpt ..\models\stable-diffusion\ +if exist "custom-model.ckpt" move custom-model.ckpt ..\models\stable-diffusion\ +if exist "GFPGANv1.3.pth" move GFPGANv1.3.pth ..\models\gfpgan\ +if exist "RealESRGAN_x4plus.pth" move RealESRGAN_x4plus.pth ..\models\realesrgan\ +if exist "RealESRGAN_x4plus_anime_6B.pth" move RealESRGAN_x4plus_anime_6B.pth ..\models\realesrgan\ + +@rem install torch and torchvision +call python ..\scripts\check_modules.py torch torchvision +if "%ERRORLEVEL%" EQU "0" ( + echo "torch and torchvision have already been installed." ) else ( - @echo. & echo "Downloading packages necessary for Stable Diffusion.." & echo. & echo "***** This will take some time (depending on the speed of the Internet connection) and may appear to be stuck, but please be patient ***** .." & echo. + echo "Installing torch and torchvision.." - @rmdir /s /q .\env + @REM prevent from using packages from the user's home directory, to avoid conflicts + set PYTHONNOUSERSITE=1 + set PYTHONPATH=%INSTALL_ENV_DIR%\lib\site-packages - @REM prevent conda from using packages from the user's home directory, to avoid conflicts - @set PYTHONNOUSERSITE=1 - - set USERPROFILE=%cd%\profile - - set PYTHONPATH=%cd%;%cd%\env\lib\site-packages - - @call conda env create --prefix env -f environment.yaml || ( - @echo. & echo "Error installing the packages necessary for Stable Diffusion. Sorry about that, please try to:" & echo " 1. Run this installer again." & echo " 2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting" & echo " 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" & echo " 4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues" & echo "Thanks!" & echo. + call python -m pip install --upgrade torch torchvision --extra-index-url https://download.pytorch.org/whl/cu116 || ( + echo "Error installing torch. Sorry about that, please try to:" & echo " 1. Run this installer again." & echo " 2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting" & echo " 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" & echo " 4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues" & echo "Thanks!" pause exit /b ) - - @call conda activate .\env - - for /f "tokens=*" %%a in ('python -c "import torch; import ldm; import transformers; import numpy; import antlr4; print(42)"') do if "%%a" NEQ "42" ( - @echo. & echo "Dependency test failed! Error installing the packages necessary for Stable Diffusion. Sorry about that, please try to:" & echo " 1. Run this installer again." & echo " 2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting" & echo " 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" & echo " 4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues" & echo "Thanks!" & echo. - pause - exit /b - ) - - @echo conda_sd_env_created >> ..\scripts\install_status.txt -) - -@rem allow rolling back the sdkit-based changes -if exist "src-old" ( - if not exist "src" ( - rename "src-old" "src" - - if exist "ldm-old" ( - rd /s /q "ldm-old" - ) - - call pip uninstall -y sdkit stable-diffusion-sdkit - ) ) set PATH=C:\Windows\System32;%PATH% -@>nul findstr /m "conda_sd_gfpgan_deps_installed" ..\scripts\install_status.txt -@if "%ERRORLEVEL%" EQU "0" ( - @echo "Packages necessary for GFPGAN (Face Correction) were already installed" +@rem install/upgrade sdkit +call python ..\scripts\check_modules.py sdkit sdkit.models ldm transformers numpy antlr4 gfpgan realesrgan +if "%ERRORLEVEL%" EQU "0" ( + echo "sdkit is already installed." + + @rem skip sdkit upgrade if in developer-mode + if not exist "..\src\sdkit" ( + @REM prevent from using packages from the user's home directory, to avoid conflicts + set PYTHONNOUSERSITE=1 + set PYTHONPATH=%INSTALL_ENV_DIR%\lib\site-packages + + call python -m pip install --upgrade sdkit -q || ( + echo "Error updating sdkit" + ) + ) ) else ( - @echo. & echo "Downloading packages necessary for GFPGAN (Face Correction).." & echo. + echo "Installing sdkit: https://pypi.org/project/sdkit/" - @set PYTHONNOUSERSITE=1 + @REM prevent from using packages from the user's home directory, to avoid conflicts + set PYTHONNOUSERSITE=1 + set PYTHONPATH=%INSTALL_ENV_DIR%\lib\site-packages - set USERPROFILE=%cd%\profile - - set PYTHONPATH=%cd%;%cd%\env\lib\site-packages - - for /f "tokens=*" %%a in ('python -c "from gfpgan import GFPGANer; print(42)"') do if "%%a" NEQ "42" ( - @echo. & echo "Dependency test failed! Error installing the packages necessary for GFPGAN (Face Correction). Sorry about that, please try to:" & echo " 1. Run this installer again." & echo " 2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting" & echo " 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" & echo " 4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues" & echo "Thanks!" & echo. + call python -m pip install sdkit || ( + echo "Error installing sdkit. Sorry about that, please try to:" & echo " 1. Run this installer again." & echo " 2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting" & echo " 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" & echo " 4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues" & echo "Thanks!" pause exit /b ) - - @echo conda_sd_gfpgan_deps_installed >> ..\scripts\install_status.txt ) -@>nul findstr /m "conda_sd_esrgan_deps_installed" ..\scripts\install_status.txt -@if "%ERRORLEVEL%" EQU "0" ( - @echo "Packages necessary for ESRGAN (Resolution Upscaling) were already installed" +call python -c "from importlib.metadata import version; print('sdkit version:', version('sdkit'))" + +@rem upgrade stable-diffusion-sdkit +call python -m pip install --upgrade stable-diffusion-sdkit -q || ( + echo "Error updating stable-diffusion-sdkit" +) +call python -c "from importlib.metadata import version; print('stable-diffusion version:', version('stable-diffusion-sdkit'))" + +@rem install rich +call python ..\scripts\check_modules.py rich +if "%ERRORLEVEL%" EQU "0" ( + echo "rich has already been installed." ) else ( - @echo. & echo "Downloading packages necessary for ESRGAN (Resolution Upscaling).." & echo. + echo "Installing rich.." - @set PYTHONNOUSERSITE=1 + set PYTHONNOUSERSITE=1 + set PYTHONPATH=%INSTALL_ENV_DIR%\lib\site-packages - set USERPROFILE=%cd%\profile - - set PYTHONPATH=%cd%;%cd%\env\lib\site-packages - - for /f "tokens=*" %%a in ('python -c "from basicsr.archs.rrdbnet_arch import RRDBNet; from realesrgan import RealESRGANer; print(42)"') do if "%%a" NEQ "42" ( - @echo. & echo "Dependency test failed! Error installing the packages necessary for ESRGAN (Resolution Upscaling). Sorry about that, please try to:" & echo " 1. Run this installer again." & echo " 2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting" & echo " 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" & echo " 4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues" & echo "Thanks!" & echo. + call python -m pip install rich || ( + echo "Error installing rich. Sorry about that, please try to:" & echo " 1. Run this installer again." & echo " 2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting" & echo " 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" & echo " 4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues" & echo "Thanks!" pause exit /b ) - - @echo conda_sd_esrgan_deps_installed >> ..\scripts\install_status.txt ) -@>nul findstr /m "conda_sd_ui_deps_installed" ..\scripts\install_status.txt +set PATH=C:\Windows\System32;%PATH% + +call python ..\scripts\check_modules.py uvicorn fastapi @if "%ERRORLEVEL%" EQU "0" ( echo "Packages necessary for Stable Diffusion UI were already installed" ) else ( @echo. & echo "Downloading packages necessary for Stable Diffusion UI.." & echo. - @set PYTHONNOUSERSITE=1 + set PYTHONNOUSERSITE=1 + set PYTHONPATH=%INSTALL_ENV_DIR%\lib\site-packages - set USERPROFILE=%cd%\profile - - set PYTHONPATH=%cd%;%cd%\env\lib\site-packages - - @call conda install -c conda-forge -y --prefix env uvicorn fastapi || ( + @call conda install -c conda-forge -y uvicorn fastapi || ( echo "Error installing the packages necessary for Stable Diffusion UI. Sorry about that, please try to:" & echo " 1. Run this installer again." & echo " 2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting" & echo " 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" & echo " 4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues" & echo "Thanks!" pause exit /b @@ -185,70 +161,35 @@ call WHERE uvicorn > .tmp exit /b ) -@>nul 2>nul call python -m picklescan --help -@if "%ERRORLEVEL%" NEQ "0" ( - @echo. & echo Picklescan not found. Installing - @call pip install picklescan || ( - echo "Error installing the picklescan package necessary for Stable Diffusion UI. Sorry about that, please try to:" & echo " 1. Run this installer again." & echo " 2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting" & echo " 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" & echo " 4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues" & echo "Thanks!" - pause - exit /b - ) -) - -@>nul 2>nul call python -c "import safetensors" -@if "%ERRORLEVEL%" NEQ "0" ( - @echo. & echo SafeTensors not found. Installing - @call pip install safetensors || ( - echo "Error installing the safetensors package necessary for Stable Diffusion UI. Sorry about that, please try to:" & echo " 1. Run this installer again." & echo " 2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting" & echo " 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" & echo " 4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues" & echo "Thanks!" - pause - exit /b - ) -) - @>nul findstr /m "conda_sd_ui_deps_installed" ..\scripts\install_status.txt @if "%ERRORLEVEL%" NEQ "0" ( @echo conda_sd_ui_deps_installed >> ..\scripts\install_status.txt ) - - -if not exist "..\models\stable-diffusion" mkdir "..\models\stable-diffusion" -if not exist "..\models\vae" mkdir "..\models\vae" -if not exist "..\models\hypernetwork" mkdir "..\models\hypernetwork" -echo. > "..\models\stable-diffusion\Put your custom ckpt files here.txt" -echo. > "..\models\vae\Put your VAE files here.txt" -echo. > "..\models\hypernetwork\Put your hypernetwork files here.txt" - -@rem if downgrading from v2.5, migrate the models to the legacy path (if already downloaded) -if exist "..\models\stable-diffusion\sd-v1-4.ckpt" if not exist "sd-v1-4.ckpt" move ..\models\stable-diffusion\sd-v1-4.ckpt . -if exist "..\models\gfpgan\GFPGANv1.3.pth" if not exist "GFPGANv1.3.pth" move ..\models\gfpgan\GFPGANv1.3.pth . -if exist "..\models\realesrgan\RealESRGAN_x4plus.pth" if not exist "RealESRGAN_x4plus.pth" move ..\models\realesrgan\RealESRGAN_x4plus.pth . -if exist "..\models\realesrgan\RealESRGAN_x4plus_anime_6B.pth" if not exist "RealESRGAN_x4plus_anime_6B.pth" move ..\models\realesrgan\RealESRGAN_x4plus_anime_6B.pth . - -@if exist "sd-v1-4.ckpt" ( - for %%I in ("sd-v1-4.ckpt") do if "%%~zI" EQU "4265380512" ( +@if exist "..\models\stable-diffusion\sd-v1-4.ckpt" ( + for %%I in ("..\models\stable-diffusion\sd-v1-4.ckpt") do if "%%~zI" EQU "4265380512" ( echo "Data files (weights) necessary for Stable Diffusion were already downloaded. Using the HuggingFace 4 GB Model." ) else ( - for %%J in ("sd-v1-4.ckpt") do if "%%~zJ" EQU "7703807346" ( + for %%J in ("..\models\stable-diffusion\sd-v1-4.ckpt") do if "%%~zJ" EQU "7703807346" ( echo "Data files (weights) necessary for Stable Diffusion were already downloaded. Using the HuggingFace 7 GB Model." ) else ( - for %%K in ("sd-v1-4.ckpt") do if "%%~zK" EQU "7703810927" ( + for %%K in ("..\models\stable-diffusion\sd-v1-4.ckpt") do if "%%~zK" EQU "7703810927" ( echo "Data files (weights) necessary for Stable Diffusion were already downloaded. Using the Waifu Model." ) else ( - echo. & echo "The model file present at %cd%\sd-v1-4.ckpt is invalid. It is only %%~zK bytes in size. Re-downloading.." & echo. - del "sd-v1-4.ckpt" + echo. & echo "The model file present at models\stable-diffusion\sd-v1-4.ckpt is invalid. It is only %%~zK bytes in size. Re-downloading.." & echo. + del "..\models\stable-diffusion\sd-v1-4.ckpt" ) ) ) ) -@if not exist "sd-v1-4.ckpt" ( +@if not exist "..\models\stable-diffusion\sd-v1-4.ckpt" ( @echo. & echo "Downloading data files (weights) for Stable Diffusion.." & echo. - @call curl -L -k https://me.cmdr2.org/stable-diffusion-ui/sd-v1-4.ckpt > sd-v1-4.ckpt + @call curl -L -k https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt > ..\models\stable-diffusion\sd-v1-4.ckpt - @if exist "sd-v1-4.ckpt" ( - for %%I in ("sd-v1-4.ckpt") do if "%%~zI" NEQ "4265380512" ( + @if exist "..\models\stable-diffusion\sd-v1-4.ckpt" ( + for %%I in ("..\models\stable-diffusion\sd-v1-4.ckpt") do if "%%~zI" NEQ "4265380512" ( echo. & echo "Error: The downloaded model file was invalid! Bytes downloaded: %%~zI" & echo. echo. & echo "Error downloading the data files (weights) for Stable Diffusion. Sorry about that, please try to:" & echo " 1. Run this installer again." & echo " 2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting" & echo " 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" & echo " 4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues" & echo "Thanks!" & echo. pause @@ -263,22 +204,22 @@ if exist "..\models\realesrgan\RealESRGAN_x4plus_anime_6B.pth" if not exist "Rea -@if exist "GFPGANv1.3.pth" ( - for %%I in ("GFPGANv1.3.pth") do if "%%~zI" EQU "348632874" ( +@if exist "..\models\gfpgan\GFPGANv1.3.pth" ( + for %%I in ("..\models\gfpgan\GFPGANv1.3.pth") do if "%%~zI" EQU "348632874" ( echo "Data files (weights) necessary for GFPGAN (Face Correction) were already downloaded" ) else ( - echo. & echo "The GFPGAN model file present at %cd%\GFPGANv1.3.pth is invalid. It is only %%~zI bytes in size. Re-downloading.." & echo. - del "GFPGANv1.3.pth" + echo. & echo "The GFPGAN model file present at models\gfpgan\GFPGANv1.3.pth is invalid. It is only %%~zI bytes in size. Re-downloading.." & echo. + del "..\models\gfpgan\GFPGANv1.3.pth" ) ) -@if not exist "GFPGANv1.3.pth" ( +@if not exist "..\models\gfpgan\GFPGANv1.3.pth" ( @echo. & echo "Downloading data files (weights) for GFPGAN (Face Correction).." & echo. - @call curl -L -k https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth > GFPGANv1.3.pth + @call curl -L -k https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth > ..\models\gfpgan\GFPGANv1.3.pth - @if exist "GFPGANv1.3.pth" ( - for %%I in ("GFPGANv1.3.pth") do if "%%~zI" NEQ "348632874" ( + @if exist "..\models\gfpgan\GFPGANv1.3.pth" ( + for %%I in ("..\models\gfpgan\GFPGANv1.3.pth") do if "%%~zI" NEQ "348632874" ( echo. & echo "Error: The downloaded GFPGAN model file was invalid! Bytes downloaded: %%~zI" & echo. echo. & echo "Error downloading the data files (weights) for GFPGAN (Face Correction). Sorry about that, please try to:" & echo " 1. Run this installer again." & echo " 2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting" & echo " 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" & echo " 4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues" & echo "Thanks!" & echo. pause @@ -293,22 +234,22 @@ if exist "..\models\realesrgan\RealESRGAN_x4plus_anime_6B.pth" if not exist "Rea -@if exist "RealESRGAN_x4plus.pth" ( - for %%I in ("RealESRGAN_x4plus.pth") do if "%%~zI" EQU "67040989" ( +@if exist "..\models\realesrgan\RealESRGAN_x4plus.pth" ( + for %%I in ("..\models\realesrgan\RealESRGAN_x4plus.pth") do if "%%~zI" EQU "67040989" ( echo "Data files (weights) necessary for ESRGAN (Resolution Upscaling) x4plus were already downloaded" ) else ( - echo. & echo "The RealESRGAN model file present at %cd%\RealESRGAN_x4plus.pth is invalid. It is only %%~zI bytes in size. Re-downloading.." & echo. - del "RealESRGAN_x4plus.pth" + echo. & echo "The RealESRGAN model file present at models\realesrgan\RealESRGAN_x4plus.pth is invalid. It is only %%~zI bytes in size. Re-downloading.." & echo. + del "..\models\realesrgan\RealESRGAN_x4plus.pth" ) ) -@if not exist "RealESRGAN_x4plus.pth" ( +@if not exist "..\models\realesrgan\RealESRGAN_x4plus.pth" ( @echo. & echo "Downloading data files (weights) for ESRGAN (Resolution Upscaling) x4plus.." & echo. - @call curl -L -k https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth > RealESRGAN_x4plus.pth + @call curl -L -k https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth > ..\models\realesrgan\RealESRGAN_x4plus.pth - @if exist "RealESRGAN_x4plus.pth" ( - for %%I in ("RealESRGAN_x4plus.pth") do if "%%~zI" NEQ "67040989" ( + @if exist "..\models\realesrgan\RealESRGAN_x4plus.pth" ( + for %%I in ("..\models\realesrgan\RealESRGAN_x4plus.pth") do if "%%~zI" NEQ "67040989" ( echo. & echo "Error: The downloaded ESRGAN x4plus model file was invalid! Bytes downloaded: %%~zI" & echo. echo. & echo "Error downloading the data files (weights) for ESRGAN (Resolution Upscaling) x4plus. Sorry about that, please try to:" & echo " 1. Run this installer again." & echo " 2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting" & echo " 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" & echo " 4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues" & echo "Thanks!" & echo. pause @@ -323,21 +264,21 @@ if exist "..\models\realesrgan\RealESRGAN_x4plus_anime_6B.pth" if not exist "Rea -@if exist "RealESRGAN_x4plus_anime_6B.pth" ( - for %%I in ("RealESRGAN_x4plus_anime_6B.pth") do if "%%~zI" EQU "17938799" ( +@if exist "..\models\realesrgan\RealESRGAN_x4plus_anime_6B.pth" ( + for %%I in ("..\models\realesrgan\RealESRGAN_x4plus_anime_6B.pth") do if "%%~zI" EQU "17938799" ( echo "Data files (weights) necessary for ESRGAN (Resolution Upscaling) x4plus_anime were already downloaded" ) else ( - echo. & echo "The RealESRGAN model file present at %cd%\RealESRGAN_x4plus_anime_6B.pth is invalid. It is only %%~zI bytes in size. Re-downloading.." & echo. - del "RealESRGAN_x4plus_anime_6B.pth" + echo. & echo "The RealESRGAN model file present at models\realesrgan\RealESRGAN_x4plus_anime_6B.pth is invalid. It is only %%~zI bytes in size. Re-downloading.." & echo. + del "..\models\realesrgan\RealESRGAN_x4plus_anime_6B.pth" ) ) -@if not exist "RealESRGAN_x4plus_anime_6B.pth" ( +@if not exist "..\models\realesrgan\RealESRGAN_x4plus_anime_6B.pth" ( @echo. & echo "Downloading data files (weights) for ESRGAN (Resolution Upscaling) x4plus_anime.." & echo. - @call curl -L -k https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth > RealESRGAN_x4plus_anime_6B.pth + @call curl -L -k https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth > ..\models\realesrgan\RealESRGAN_x4plus_anime_6B.pth - @if exist "RealESRGAN_x4plus_anime_6B.pth" ( + @if exist "..\models\realesrgan\RealESRGAN_x4plus_anime_6B.pth" ( for %%I in ("RealESRGAN_x4plus_anime_6B.pth") do if "%%~zI" NEQ "17938799" ( echo. & echo "Error: The downloaded ESRGAN x4plus_anime model file was invalid! Bytes downloaded: %%~zI" & echo. echo. & echo "Error downloading the data files (weights) for ESRGAN (Resolution Upscaling) x4plus_anime. Sorry about that, please try to:" & echo " 1. Run this installer again." & echo " 2. If that doesn't fix it, please try the common troubleshooting steps at https://github.com/cmdr2/stable-diffusion-ui/wiki/Troubleshooting" & echo " 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" & echo " 4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues" & echo "Thanks!" & echo. @@ -381,10 +322,6 @@ if exist "..\models\realesrgan\RealESRGAN_x4plus_anime_6B.pth" if not exist "Rea ) ) -if "%test_sd2%" == "Y" ( - @call pip install open_clip_torch==2.0.2 -) - @>nul findstr /m "sd_install_complete" ..\scripts\install_status.txt @if "%ERRORLEVEL%" NEQ "0" ( @echo sd_weights_downloaded >> ..\scripts\install_status.txt @@ -395,10 +332,8 @@ if "%test_sd2%" == "Y" ( @set SD_DIR=%cd% -@cd env\lib\site-packages -@set PYTHONPATH=%SD_DIR%;%cd% -@cd ..\..\.. -@echo PYTHONPATH=%PYTHONPATH% +set PYTHONPATH=%INSTALL_ENV_DIR%\lib\site-packages +echo PYTHONPATH=%PYTHONPATH% call where python call python --version @@ -407,17 +342,12 @@ call python --version @set SD_UI_PATH=%cd%\ui @cd stable-diffusion -@rem -@rem Rewrite easy-install.pth. This fixes the installation if the user has relocated the SDUI installation -@rem ->env\Lib\site-packages\easy-install.pth echo %cd%\src\taming-transformers ->>env\Lib\site-packages\easy-install.pth echo %cd%\src\clip ->>env\Lib\site-packages\easy-install.pth echo %cd%\src\gfpgan ->>env\Lib\site-packages\easy-install.pth echo %cd%\src\realesrgan +@rem set any overrides +set HF_HUB_DISABLE_SYMLINKS_WARNING=true @if NOT DEFINED SD_UI_BIND_PORT set SD_UI_BIND_PORT=9000 @if NOT DEFINED SD_UI_BIND_IP set SD_UI_BIND_IP=0.0.0.0 -@uvicorn server:app --app-dir "%SD_UI_PATH%" --port %SD_UI_BIND_PORT% --host %SD_UI_BIND_IP% +@uvicorn main:server_api --app-dir "%SD_UI_PATH%" --port %SD_UI_BIND_PORT% --host %SD_UI_BIND_IP% --log-level error @pause diff --git a/scripts/on_sd_start.sh b/scripts/on_sd_start.sh index 44addb7e..c15af743 100755 --- a/scripts/on_sd_start.sh +++ b/scripts/on_sd_start.sh @@ -4,6 +4,7 @@ source ./scripts/functions.sh cp sd-ui-files/scripts/on_env_start.sh scripts/ cp sd-ui-files/scripts/bootstrap.sh scripts/ +cp sd-ui-files/scripts/check_modules.py scripts/ # activate the installer env CONDA_BASEPATH=$(conda info --base) @@ -21,125 +22,110 @@ python -c "import os; import shutil; frm = 'sd-ui-files/ui/hotfix/9c24e6cd9f499d # Caution, this file will make your eyes and brain bleed. It's such an unholy mess. # Note to self: Please rewrite this in Python. For the sake of your own sanity. -if [ "$test_sd2" == "" ]; then - export test_sd2="N" -fi - -if [ -e "scripts/install_status.txt" ] && [ `grep -c sd_git_cloned scripts/install_status.txt` -gt "0" ]; then - echo "Stable Diffusion's git repository was already installed. Updating.." - - cd stable-diffusion - - git remote set-url origin https://github.com/easydiffusion/diffusion-kit.git - - git reset --hard - git pull - - if [ "$test_sd2" == "N" ]; then - git -c advice.detachedHead=false checkout 7f32368ed1030a6e710537047bacd908adea183a - elif [ "$test_sd2" == "Y" ]; then - git -c advice.detachedHead=false checkout 733a1f6f9cae9b9a9b83294bf3281b123378cb1f - fi - - cd .. -else - printf "\n\nDownloading Stable Diffusion..\n\n" - - if git clone https://github.com/easydiffusion/diffusion-kit.git stable-diffusion ; then - echo sd_git_cloned >> scripts/install_status.txt - else - fail "git clone of basujindal/stable-diffusion.git failed" - fi - - cd stable-diffusion - git -c advice.detachedHead=false checkout 7f32368ed1030a6e710537047bacd908adea183a - - cd .. +# set the correct installer path (current vs legacy) +if [ -e "installer_files/env" ]; then + export INSTALL_ENV_DIR="$(pwd)/installer_files/env" +fi +if [ -e "stable-diffusion/env" ]; then + export INSTALL_ENV_DIR="$(pwd)/stable-diffusion/env" fi +# create the stable-diffusion folder, to work with legacy installations +if [ ! -e "stable-diffusion" ]; then mkdir stable-diffusion; fi cd stable-diffusion -if [ `grep -c conda_sd_env_created ../scripts/install_status.txt` -gt "0" ]; then - echo "Packages necessary for Stable Diffusion were already installed" - +# activate the old stable-diffusion env, if it exists +if [ -e "env" ]; then conda activate ./env || fail "conda activate failed" +fi + +# disable the legacy src and ldm folder (otherwise this prevents installing gfpgan and realesrgan) +if [ -e "src" ]; then mv src src-old; fi +if [ -e "ldm" ]; then mv ldm ldm-old; fi + +mkdir -p "../models/stable-diffusion" +mkdir -p "../models/gfpgan" +mkdir -p "../models/realesrgan" +mkdir -p "../models/vae" + +# migrate the legacy models to the correct path (if already downloaded) +if [ -e "sd-v1-4.ckpt" ]; then mv sd-v1-4.ckpt ../models/stable-diffusion/; fi +if [ -e "custom-model.ckpt" ]; then mv custom-model.ckpt ../models/stable-diffusion/; fi +if [ -e "GFPGANv1.3.pth" ]; then mv GFPGANv1.3.pth ../models/gfpgan/; fi +if [ -e "RealESRGAN_x4plus.pth" ]; then mv RealESRGAN_x4plus.pth ../models/realesrgan/; fi +if [ -e "RealESRGAN_x4plus_anime_6B.pth" ]; then mv RealESRGAN_x4plus_anime_6B.pth ../models/realesrgan/; fi + +# install torch and torchvision +if python ../scripts/check_modules.py torch torchvision; then + echo "torch and torchvision have already been installed." else - printf "\n\nDownloading packages necessary for Stable Diffusion..\n" - printf "\n\n***** This will take some time (depending on the speed of the Internet connection) and may appear to be stuck, but please be patient ***** ..\n\n" + echo "Installing torch and torchvision.." - # prevent conda from using packages from the user's home directory, to avoid conflicts export PYTHONNOUSERSITE=1 - export PYTHONPATH="$(pwd):$(pwd)/env/lib/site-packages" + export PYTHONPATH="$INSTALL_ENV_DIR/lib/python3.8/site-packages" - if conda env create --prefix env --force -f environment.yaml ; then - echo "Installed. Testing.." + if python -m pip install --upgrade torch torchvision --extra-index-url https://download.pytorch.org/whl/cu116 ; then + echo "Installed." else - fail "'conda env create' failed" + fail "torch install failed" fi - - conda activate ./env || fail "conda activate failed" - - out_test=`python -c "import torch; import ldm; import transformers; import numpy; import antlr4; print(42)"` - if [ "$out_test" != "42" ]; then - fail "Dependency test failed" - fi - - echo conda_sd_env_created >> ../scripts/install_status.txt fi -# allow rolling back the sdkit-based changes -if [ -e "src-old" ] && [ ! -e "src" ]; then - mv src-old src +# install/upgrade sdkit +if python ../scripts/check_modules.py sdkit sdkit.models ldm transformers numpy antlr4 gfpgan realesrgan ; then + echo "sdkit is already installed." - if [ -e "ldm-old" ]; then rm -r ldm-old; fi + # skip sdkit upgrade if in developer-mode + if [ ! -e "../src/sdkit" ]; then + export PYTHONNOUSERSITE=1 + export PYTHONPATH="$INSTALL_ENV_DIR/lib/python3.8/site-packages" - pip uninstall -y sdkit stable-diffusion-sdkit -fi - -if [ `grep -c conda_sd_gfpgan_deps_installed ../scripts/install_status.txt` -gt "0" ]; then - echo "Packages necessary for GFPGAN (Face Correction) were already installed" + python -m pip install --upgrade sdkit -q + fi else - printf "\n\nDownloading packages necessary for GFPGAN (Face Correction)..\n" + echo "Installing sdkit: https://pypi.org/project/sdkit/" export PYTHONNOUSERSITE=1 - export PYTHONPATH="$(pwd):$(pwd)/env/lib/site-packages" + export PYTHONPATH="$INSTALL_ENV_DIR/lib/python3.8/site-packages" - out_test=`python -c "from gfpgan import GFPGANer; print(42)"` - if [ "$out_test" != "42" ]; then - echo "EE The dependency check has failed. This usually means that some system libraries are missing." - echo "EE On Debian/Ubuntu systems, this are often these packages: libsm6 libxext6 libxrender-dev" - echo "EE Other Linux distributions might have different package names for these libraries." - fail "GFPGAN dependency test failed" + if python -m pip install sdkit ; then + echo "Installed." + else + fail "sdkit install failed" fi - - echo conda_sd_gfpgan_deps_installed >> ../scripts/install_status.txt fi -if [ `grep -c conda_sd_esrgan_deps_installed ../scripts/install_status.txt` -gt "0" ]; then - echo "Packages necessary for ESRGAN (Resolution Upscaling) were already installed" +python -c "from importlib.metadata import version; print('sdkit version:', version('sdkit'))" + +# upgrade stable-diffusion-sdkit +python -m pip install --upgrade stable-diffusion-sdkit -q +python -c "from importlib.metadata import version; print('stable-diffusion version:', version('stable-diffusion-sdkit'))" + +# install rich +if python ../scripts/check_modules.py rich; then + echo "rich has already been installed." else - printf "\n\nDownloading packages necessary for ESRGAN (Resolution Upscaling)..\n" + echo "Installing rich.." export PYTHONNOUSERSITE=1 - export PYTHONPATH="$(pwd):$(pwd)/env/lib/site-packages" + export PYTHONPATH="$INSTALL_ENV_DIR/lib/python3.8/site-packages" - out_test=`python -c "from basicsr.archs.rrdbnet_arch import RRDBNet; from realesrgan import RealESRGANer; print(42)"` - if [ "$out_test" != "42" ]; then - fail "ESRGAN dependency test failed" + if python -m pip install rich ; then + echo "Installed." + else + fail "Install failed for rich" fi - - echo conda_sd_esrgan_deps_installed >> ../scripts/install_status.txt fi -if [ `grep -c conda_sd_ui_deps_installed ../scripts/install_status.txt` -gt "0" ]; then +if python ../scripts/check_modules.py uvicorn fastapi ; then echo "Packages necessary for Stable Diffusion UI were already installed" else printf "\n\nDownloading packages necessary for Stable Diffusion UI..\n\n" export PYTHONNOUSERSITE=1 - export PYTHONPATH="$(pwd):$(pwd)/env/lib/site-packages" + export PYTHONPATH="$INSTALL_ENV_DIR/lib/python3.8/site-packages" - if conda install -c conda-forge --prefix ./env -y uvicorn fastapi ; then + if conda install -c conda-forge -y uvicorn fastapi ; then echo "Installed. Testing.." else fail "'conda install uvicorn' failed" @@ -148,57 +134,26 @@ else if ! command -v uvicorn &> /dev/null; then fail "UI packages not found!" fi - - echo conda_sd_ui_deps_installed >> ../scripts/install_status.txt fi -if python -m picklescan --help >/dev/null 2>&1; then - echo "Picklescan is already installed." -else - echo "Picklescan not found, installing." - pip install picklescan || fail "Picklescan installation failed." -fi - -if python -c "import safetensors" --help >/dev/null 2>&1; then - echo "SafeTensors is already installed." -else - echo "SafeTensors not found, installing." - pip install safetensors || fail "SafeTensors installation failed." -fi - - - -mkdir -p "../models/stable-diffusion" -mkdir -p "../models/vae" -mkdir -p "../models/hypernetwork" -echo "" > "../models/stable-diffusion/Put your custom ckpt files here.txt" -echo "" > "../models/vae/Put your VAE files here.txt" -echo "" > "../models/hypernetwork/Put your hypernetwork files here.txt" - -# if downgrading from v2.5, migrate the models to the legacy path (if already downloaded) -if [ -e "../models/stable-diffusion/sd-v1-4.ckpt" ] && [ ! -e "sd-v1-4.ckpt" ]; then mv ../models/stable-diffusion/sd-v1-4.ckpt . ; fi -if [ -e "../models/gfpgan/GFPGANv1.3.pth" ] && [ ! -e "GFPGANv1.3.pth" ]; then mv ../models/gfpgan/GFPGANv1.3.pth . ; fi -if [ -e "../models/realesrgan/RealESRGAN_x4plus.pth" ] && [ ! -e "RealESRGAN_x4plus.pth" ]; then mv ../models/realesrgan/RealESRGAN_x4plus.pth . ; fi -if [ -e "../models/realesrgan/RealESRGAN_x4plus_anime_6B.pth" ] && [ ! -e "RealESRGAN_x4plus_anime_6B.pth" ]; then mv ../models/realesrgan/RealESRGAN_x4plus_anime_6B.pth . ; fi - -if [ -f "sd-v1-4.ckpt" ]; then - model_size=`find "sd-v1-4.ckpt" -printf "%s"` +if [ -f "../models/stable-diffusion/sd-v1-4.ckpt" ]; then + model_size=`find "../models/stable-diffusion/sd-v1-4.ckpt" -printf "%s"` if [ "$model_size" -eq "4265380512" ] || [ "$model_size" -eq "7703807346" ] || [ "$model_size" -eq "7703810927" ]; then echo "Data files (weights) necessary for Stable Diffusion were already downloaded" else - printf "\n\nThe model file present at $PWD/sd-v1-4.ckpt is invalid. It is only $model_size bytes in size. Re-downloading.." - rm sd-v1-4.ckpt + printf "\n\nThe model file present at models/stable-diffusion/sd-v1-4.ckpt is invalid. It is only $model_size bytes in size. Re-downloading.." + rm ../models/stable-diffusion/sd-v1-4.ckpt fi fi -if [ ! -f "sd-v1-4.ckpt" ]; then +if [ ! -f "../models/stable-diffusion/sd-v1-4.ckpt" ]; then echo "Downloading data files (weights) for Stable Diffusion.." - curl -L -k https://me.cmdr2.org/stable-diffusion-ui/sd-v1-4.ckpt > sd-v1-4.ckpt + curl -L -k https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt > ../models/stable-diffusion/sd-v1-4.ckpt - if [ -f "sd-v1-4.ckpt" ]; then - model_size=`find "sd-v1-4.ckpt" -printf "%s"` + if [ -f "../models/stable-diffusion/sd-v1-4.ckpt" ]; then + model_size=`find "../models/stable-diffusion/sd-v1-4.ckpt" -printf "%s"` if [ ! "$model_size" == "4265380512" ]; then fail "The downloaded model file was invalid! Bytes downloaded: $model_size" fi @@ -208,24 +163,24 @@ if [ ! -f "sd-v1-4.ckpt" ]; then fi -if [ -f "GFPGANv1.3.pth" ]; then - model_size=`find "GFPGANv1.3.pth" -printf "%s"` +if [ -f "../models/gfpgan/GFPGANv1.3.pth" ]; then + model_size=`find "../models/gfpgan/GFPGANv1.3.pth" -printf "%s"` if [ "$model_size" -eq "348632874" ]; then echo "Data files (weights) necessary for GFPGAN (Face Correction) were already downloaded" else - printf "\n\nThe model file present at $PWD/GFPGANv1.3.pth is invalid. It is only $model_size bytes in size. Re-downloading.." - rm GFPGANv1.3.pth + printf "\n\nThe model file present at models/gfpgan/GFPGANv1.3.pth is invalid. It is only $model_size bytes in size. Re-downloading.." + rm ../models/gfpgan/GFPGANv1.3.pth fi fi -if [ ! -f "GFPGANv1.3.pth" ]; then +if [ ! -f "../models/gfpgan/GFPGANv1.3.pth" ]; then echo "Downloading data files (weights) for GFPGAN (Face Correction).." - curl -L -k https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth > GFPGANv1.3.pth + curl -L -k https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth > ../models/gfpgan/GFPGANv1.3.pth - if [ -f "GFPGANv1.3.pth" ]; then - model_size=`find "GFPGANv1.3.pth" -printf "%s"` + if [ -f "../models/gfpgan/GFPGANv1.3.pth" ]; then + model_size=`find "../models/gfpgan/GFPGANv1.3.pth" -printf "%s"` if [ ! "$model_size" -eq "348632874" ]; then fail "The downloaded GFPGAN model file was invalid! Bytes downloaded: $model_size" fi @@ -235,24 +190,24 @@ if [ ! -f "GFPGANv1.3.pth" ]; then fi -if [ -f "RealESRGAN_x4plus.pth" ]; then - model_size=`find "RealESRGAN_x4plus.pth" -printf "%s"` +if [ -f "../models/realesrgan/RealESRGAN_x4plus.pth" ]; then + model_size=`find "../models/realesrgan/RealESRGAN_x4plus.pth" -printf "%s"` if [ "$model_size" -eq "67040989" ]; then echo "Data files (weights) necessary for ESRGAN (Resolution Upscaling) x4plus were already downloaded" else - printf "\n\nThe model file present at $PWD/RealESRGAN_x4plus.pth is invalid. It is only $model_size bytes in size. Re-downloading.." - rm RealESRGAN_x4plus.pth + printf "\n\nThe model file present at models/realesrgan/RealESRGAN_x4plus.pth is invalid. It is only $model_size bytes in size. Re-downloading.." + rm ../models/realesrgan/RealESRGAN_x4plus.pth fi fi -if [ ! -f "RealESRGAN_x4plus.pth" ]; then +if [ ! -f "../models/realesrgan/RealESRGAN_x4plus.pth" ]; then echo "Downloading data files (weights) for ESRGAN (Resolution Upscaling) x4plus.." - curl -L -k https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth > RealESRGAN_x4plus.pth + curl -L -k https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth > ../models/realesrgan/RealESRGAN_x4plus.pth - if [ -f "RealESRGAN_x4plus.pth" ]; then - model_size=`find "RealESRGAN_x4plus.pth" -printf "%s"` + if [ -f "../models/realesrgan/RealESRGAN_x4plus.pth" ]; then + model_size=`find "../models/realesrgan/RealESRGAN_x4plus.pth" -printf "%s"` if [ ! "$model_size" -eq "67040989" ]; then fail "The downloaded ESRGAN x4plus model file was invalid! Bytes downloaded: $model_size" fi @@ -262,24 +217,24 @@ if [ ! -f "RealESRGAN_x4plus.pth" ]; then fi -if [ -f "RealESRGAN_x4plus_anime_6B.pth" ]; then - model_size=`find "RealESRGAN_x4plus_anime_6B.pth" -printf "%s"` +if [ -f "../models/realesrgan/RealESRGAN_x4plus_anime_6B.pth" ]; then + model_size=`find "../models/realesrgan/RealESRGAN_x4plus_anime_6B.pth" -printf "%s"` if [ "$model_size" -eq "17938799" ]; then echo "Data files (weights) necessary for ESRGAN (Resolution Upscaling) x4plus_anime were already downloaded" else - printf "\n\nThe model file present at $PWD/RealESRGAN_x4plus_anime_6B.pth is invalid. It is only $model_size bytes in size. Re-downloading.." - rm RealESRGAN_x4plus_anime_6B.pth + printf "\n\nThe model file present at models/realesrgan/RealESRGAN_x4plus_anime_6B.pth is invalid. It is only $model_size bytes in size. Re-downloading.." + rm ../models/realesrgan/RealESRGAN_x4plus_anime_6B.pth fi fi -if [ ! -f "RealESRGAN_x4plus_anime_6B.pth" ]; then +if [ ! -f "../models/realesrgan/RealESRGAN_x4plus_anime_6B.pth" ]; then echo "Downloading data files (weights) for ESRGAN (Resolution Upscaling) x4plus_anime.." - curl -L -k https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth > RealESRGAN_x4plus_anime_6B.pth + curl -L -k https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth > ../models/realesrgan/RealESRGAN_x4plus_anime_6B.pth - if [ -f "RealESRGAN_x4plus_anime_6B.pth" ]; then - model_size=`find "RealESRGAN_x4plus_anime_6B.pth" -printf "%s"` + if [ -f "../models/realesrgan/RealESRGAN_x4plus_anime_6B.pth" ]; then + model_size=`find "../models/realesrgan/RealESRGAN_x4plus_anime_6B.pth" -printf "%s"` if [ ! "$model_size" -eq "17938799" ]; then fail "The downloaded ESRGAN x4plus_anime model file was invalid! Bytes downloaded: $model_size" fi @@ -320,10 +275,6 @@ if [ ! -f "../models/vae/vae-ft-mse-840000-ema-pruned.ckpt" ]; then fi fi -if [ "$test_sd2" == "Y" ]; then - pip install open_clip_torch==2.0.2 -fi - if [ `grep -c sd_install_complete ../scripts/install_status.txt` -gt "0" ]; then echo sd_weights_downloaded >> ../scripts/install_status.txt echo sd_install_complete >> ../scripts/install_status.txt @@ -332,7 +283,8 @@ fi printf "\n\nStable Diffusion is ready!\n\n" SD_PATH=`pwd` -export PYTHONPATH="$SD_PATH:$SD_PATH/env/lib/python3.8/site-packages" + +export PYTHONPATH="$INSTALL_ENV_DIR/lib/python3.8/site-packages" echo "PYTHONPATH=$PYTHONPATH" which python @@ -342,6 +294,6 @@ cd .. export SD_UI_PATH=`pwd`/ui cd stable-diffusion -uvicorn server:app --app-dir "$SD_UI_PATH" --port ${SD_UI_BIND_PORT:-9000} --host ${SD_UI_BIND_IP:-0.0.0.0} +uvicorn main:server_api --app-dir "$SD_UI_PATH" --port ${SD_UI_BIND_PORT:-9000} --host ${SD_UI_BIND_IP:-0.0.0.0} --log-level error read -p "Press any key to continue" diff --git a/scripts/start.sh b/scripts/start.sh index 4bb6cb66..61ac3f74 100755 --- a/scripts/start.sh +++ b/scripts/start.sh @@ -19,4 +19,5 @@ which conda conda --version || exit 1 # Download the rest of the installer and UI +chmod +x scripts/*.sh scripts/on_env_start.sh diff --git a/ui/easydiffusion/__init__.py b/ui/easydiffusion/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ui/easydiffusion/app.py b/ui/easydiffusion/app.py new file mode 100644 index 00000000..63505ed3 --- /dev/null +++ b/ui/easydiffusion/app.py @@ -0,0 +1,165 @@ +import os +import socket +import sys +import json +import traceback +import logging +from rich.logging import RichHandler + +from sdkit.utils import log as sdkit_log # hack, so we can overwrite the log config + +from easydiffusion import task_manager +from easydiffusion.utils import log + +# Remove all handlers associated with the root logger object. +for handler in logging.root.handlers[:]: + logging.root.removeHandler(handler) + +LOG_FORMAT = '%(asctime)s.%(msecs)03d %(levelname)s %(threadName)s %(message)s' +logging.basicConfig( + level=logging.INFO, + format=LOG_FORMAT, + datefmt="%X", + handlers=[RichHandler(markup=True, rich_tracebacks=False, show_time=False, show_level=False)], +) + +SD_DIR = os.getcwd() + +SD_UI_DIR = os.getenv('SD_UI_PATH', None) +sys.path.append(os.path.dirname(SD_UI_DIR)) + +CONFIG_DIR = os.path.abspath(os.path.join(SD_UI_DIR, '..', 'scripts')) +MODELS_DIR = os.path.abspath(os.path.join(SD_DIR, '..', 'models')) + +USER_UI_PLUGINS_DIR = os.path.abspath(os.path.join(SD_DIR, '..', 'plugins', 'ui')) +CORE_UI_PLUGINS_DIR = os.path.abspath(os.path.join(SD_UI_DIR, 'plugins', 'ui')) +UI_PLUGINS_SOURCES = ((CORE_UI_PLUGINS_DIR, 'core'), (USER_UI_PLUGINS_DIR, 'user')) + +OUTPUT_DIRNAME = "Stable Diffusion UI" # in the user's home folder +TASK_TTL = 15 * 60 # Discard last session's task timeout +APP_CONFIG_DEFAULTS = { + # auto: selects the cuda device with the most free memory, cuda: use the currently active cuda device. + 'render_devices': 'auto', # valid entries: 'auto', 'cpu' or 'cuda:N' (where N is a GPU index) + 'update_branch': 'main', + 'ui': { + 'open_browser_on_start': True, + }, +} + +def init(): + os.makedirs(USER_UI_PLUGINS_DIR, exist_ok=True) + + update_render_threads() + +def getConfig(default_val=APP_CONFIG_DEFAULTS): + try: + config_json_path = os.path.join(CONFIG_DIR, 'config.json') + if not os.path.exists(config_json_path): + return default_val + with open(config_json_path, 'r', encoding='utf-8') as f: + config = json.load(f) + if 'net' not in config: + config['net'] = {} + if os.getenv('SD_UI_BIND_PORT') is not None: + config['net']['listen_port'] = int(os.getenv('SD_UI_BIND_PORT')) + if os.getenv('SD_UI_BIND_IP') is not None: + config['net']['listen_to_network'] = (os.getenv('SD_UI_BIND_IP') == '0.0.0.0') + return config + except Exception as e: + log.warn(traceback.format_exc()) + return default_val + +def setConfig(config): + try: # config.json + config_json_path = os.path.join(CONFIG_DIR, 'config.json') + with open(config_json_path, 'w', encoding='utf-8') as f: + json.dump(config, f) + except: + log.error(traceback.format_exc()) + + try: # config.bat + config_bat_path = os.path.join(CONFIG_DIR, 'config.bat') + config_bat = [] + + if 'update_branch' in config: + config_bat.append(f"@set update_branch={config['update_branch']}") + + config_bat.append(f"@set SD_UI_BIND_PORT={config['net']['listen_port']}") + bind_ip = '0.0.0.0' if config['net']['listen_to_network'] else '127.0.0.1' + config_bat.append(f"@set SD_UI_BIND_IP={bind_ip}") + + if len(config_bat) > 0: + with open(config_bat_path, 'w', encoding='utf-8') as f: + f.write('\r\n'.join(config_bat)) + except: + log.error(traceback.format_exc()) + + try: # config.sh + config_sh_path = os.path.join(CONFIG_DIR, 'config.sh') + config_sh = ['#!/bin/bash'] + + if 'update_branch' in config: + config_sh.append(f"export update_branch={config['update_branch']}") + + config_sh.append(f"export SD_UI_BIND_PORT={config['net']['listen_port']}") + bind_ip = '0.0.0.0' if config['net']['listen_to_network'] else '127.0.0.1' + config_sh.append(f"export SD_UI_BIND_IP={bind_ip}") + + if len(config_sh) > 1: + with open(config_sh_path, 'w', encoding='utf-8') as f: + f.write('\n'.join(config_sh)) + except: + log.error(traceback.format_exc()) + +def save_to_config(ckpt_model_name, vae_model_name, hypernetwork_model_name, vram_usage_level): + config = getConfig() + if 'model' not in config: + config['model'] = {} + + config['model']['stable-diffusion'] = ckpt_model_name + config['model']['vae'] = vae_model_name + config['model']['hypernetwork'] = hypernetwork_model_name + + if vae_model_name is None or vae_model_name == "": + del config['model']['vae'] + if hypernetwork_model_name is None or hypernetwork_model_name == "": + del config['model']['hypernetwork'] + + config['vram_usage_level'] = vram_usage_level + + setConfig(config) + +def update_render_threads(): + config = getConfig() + render_devices = config.get('render_devices', 'auto') + active_devices = task_manager.get_devices()['active'].keys() + + log.debug(f'requesting for render_devices: {render_devices}') + task_manager.update_render_threads(render_devices, active_devices) + +def getUIPlugins(): + plugins = [] + + for plugins_dir, dir_prefix in UI_PLUGINS_SOURCES: + for file in os.listdir(plugins_dir): + if file.endswith('.plugin.js'): + plugins.append(f'/plugins/{dir_prefix}/{file}') + + return plugins + +def getIPConfig(): + try: + ips = socket.gethostbyname_ex(socket.gethostname()) + ips[2].append(ips[0]) + return ips[2] + except Exception as e: + log.exception(e) + return [] + +def open_browser(): + config = getConfig() + ui = config.get('ui', {}) + net = config.get('net', {'listen_port':9000}) + port = net.get('listen_port', 9000) + if ui.get('open_browser_on_start', True): + import webbrowser; webbrowser.open(f"http://localhost:{port}") diff --git a/ui/sd_internal/device_manager.py b/ui/easydiffusion/device_manager.py similarity index 69% rename from ui/sd_internal/device_manager.py rename to ui/easydiffusion/device_manager.py index d2c6430b..77fcbaea 100644 --- a/ui/sd_internal/device_manager.py +++ b/ui/easydiffusion/device_manager.py @@ -3,6 +3,15 @@ import torch import traceback import re +from easydiffusion.utils import log + +''' +Set `FORCE_FULL_PRECISION` in the environment variables, or in `config.bat`/`config.sh` to set full precision (i.e. float32). +Otherwise the models will load at half-precision (i.e. float16). + +Half-precision is fine most of the time. Full precision is only needed for working around GPU bugs (like NVIDIA 16xx GPUs). +''' + COMPARABLE_GPU_PERCENTILE = 0.65 # if a GPU's free_mem is within this % of the GPU with the most free_mem, it will be picked mem_free_threshold = 0 @@ -34,7 +43,7 @@ def get_device_delta(render_devices, active_devices): if 'auto' in render_devices: render_devices = auto_pick_devices(active_devices) if 'cpu' in render_devices: - print('WARNING: Could not find a compatible GPU. Using the CPU, but this will be very slow!') + log.warn('WARNING: Could not find a compatible GPU. Using the CPU, but this will be very slow!') active_devices = set(active_devices) render_devices = set(render_devices) @@ -53,7 +62,7 @@ def auto_pick_devices(currently_active_devices): if device_count == 1: return ['cuda:0'] if is_device_compatible('cuda:0') else ['cpu'] - print('Autoselecting GPU. Using most free memory.') + log.debug('Autoselecting GPU. Using most free memory.') devices = [] for device in range(device_count): device = f'cuda:{device}' @@ -64,7 +73,7 @@ def auto_pick_devices(currently_active_devices): mem_free /= float(10**9) mem_total /= float(10**9) device_name = torch.cuda.get_device_name(device) - print(f'{device} detected: {device_name} - Memory (free/total): {round(mem_free, 2)}Gb / {round(mem_total, 2)}Gb') + log.debug(f'{device} detected: {device_name} - Memory (free/total): {round(mem_free, 2)}Gb / {round(mem_total, 2)}Gb') devices.append({'device': device, 'device_name': device_name, 'mem_free': mem_free}) devices.sort(key=lambda x:x['mem_free'], reverse=True) @@ -82,7 +91,7 @@ def auto_pick_devices(currently_active_devices): devices = list(map(lambda x: x['device'], devices)) return devices -def device_init(thread_data, device): +def device_init(context, device): ''' This function assumes the 'device' has already been verified to be compatible. `get_device_delta()` has already filtered out incompatible devices. @@ -91,27 +100,45 @@ def device_init(thread_data, device): validate_device_id(device, log_prefix='device_init') if device == 'cpu': - thread_data.device = 'cpu' - thread_data.device_name = get_processor_name() - print('Render device CPU available as', thread_data.device_name) + context.device = 'cpu' + context.device_name = get_processor_name() + context.half_precision = False + log.debug(f'Render device CPU available as {context.device_name}') return - thread_data.device_name = torch.cuda.get_device_name(device) - thread_data.device = device + context.device_name = torch.cuda.get_device_name(device) + context.device = device # Force full precision on 1660 and 1650 NVIDIA cards to avoid creating green images - device_name = thread_data.device_name.lower() - thread_data.force_full_precision = (('nvidia' in device_name or 'geforce' in device_name) and (' 1660' in device_name or ' 1650' in device_name)) or ('Quadro T2000' in device_name) - if thread_data.force_full_precision: - print('forcing full precision on NVIDIA 16xx cards, to avoid green images. GPU detected: ', thread_data.device_name) + if needs_to_force_full_precision(context): + log.warn(f'forcing full precision on this GPU, to avoid green images. GPU detected: {context.device_name}') # Apply force_full_precision now before models are loaded. - thread_data.precision = 'full' + context.half_precision = False - print(f'Setting {device} as active') + log.info(f'Setting {device} as active, with precision: {"half" if context.half_precision else "full"}') torch.cuda.device(device) return +def needs_to_force_full_precision(context): + if 'FORCE_FULL_PRECISION' in os.environ: + return True + + device_name = context.device_name.lower() + return (('nvidia' in device_name or 'geforce' in device_name or 'quadro' in device_name) and (' 1660' in device_name or ' 1650' in device_name or ' t400' in device_name or ' t500' in device_name or ' t550' in device_name or ' t600' in device_name or ' t1000' in device_name or ' t1200' in device_name or ' t2000' in device_name)) + +def get_max_vram_usage_level(device): + if device != 'cpu': + _, mem_total = torch.cuda.mem_get_info(device) + mem_total /= float(10**9) + + if mem_total < 4.5: + return 'low' + elif mem_total < 6.5: + return 'balanced' + + return 'high' + def validate_device_id(device, log_prefix=''): def is_valid(): if not isinstance(device, str): @@ -129,10 +156,12 @@ def is_device_compatible(device): ''' Returns True/False, and prints any compatibility errors ''' + # static variable "history". + is_device_compatible.history = getattr(is_device_compatible, 'history', {}) try: validate_device_id(device, log_prefix='is_device_compatible') except: - print(str(e)) + log.error(str(e)) return False if device == 'cpu': return True @@ -141,10 +170,12 @@ def is_device_compatible(device): _, mem_total = torch.cuda.mem_get_info(device) mem_total /= float(10**9) if mem_total < 3.0: - print(f'GPU {device} with less than 3 GB of VRAM is not compatible with Stable Diffusion') + if is_device_compatible.history.get(device) == None: + log.warn(f'GPU {device} with less than 3 GB of VRAM is not compatible with Stable Diffusion') + is_device_compatible.history[device] = 1 return False except RuntimeError as e: - print(str(e)) + log.error(str(e)) return False return True @@ -164,5 +195,5 @@ def get_processor_name(): if "model name" in line: return re.sub(".*model name.*:", "", line, 1).strip() except: - print(traceback.format_exc()) + log.error(traceback.format_exc()) return "cpu" diff --git a/ui/easydiffusion/model_manager.py b/ui/easydiffusion/model_manager.py new file mode 100644 index 00000000..70bcdacf --- /dev/null +++ b/ui/easydiffusion/model_manager.py @@ -0,0 +1,222 @@ +import os + +from easydiffusion import app, device_manager +from easydiffusion.types import TaskData +from easydiffusion.utils import log + +from sdkit import Context +from sdkit.models import load_model, unload_model, get_model_info_from_db, scan_model +from sdkit.utils import hash_file_quick + +KNOWN_MODEL_TYPES = ['stable-diffusion', 'vae', 'hypernetwork', 'gfpgan', 'realesrgan'] +MODEL_EXTENSIONS = { + 'stable-diffusion': ['.ckpt', '.safetensors'], + 'vae': ['.vae.pt', '.ckpt', '.safetensors'], + 'hypernetwork': ['.pt', '.safetensors'], + 'gfpgan': ['.pth'], + 'realesrgan': ['.pth'], +} +DEFAULT_MODELS = { + 'stable-diffusion': [ # needed to support the legacy installations + 'custom-model', # only one custom model file was supported initially, creatively named 'custom-model' + 'sd-v1-4', # Default fallback. + ], + 'gfpgan': ['GFPGANv1.3'], + 'realesrgan': ['RealESRGAN_x4plus'], +} +MODELS_TO_LOAD_ON_START = ['stable-diffusion', 'vae', 'hypernetwork'] + +known_models = {} + +def init(): + make_model_folders() + getModels() # run this once, to cache the picklescan results + +def load_default_models(context: Context): + set_vram_optimizations(context) + + # init default model paths + for model_type in MODELS_TO_LOAD_ON_START: + context.model_paths[model_type] = resolve_model_to_use(model_type=model_type) + try: + load_model(context, model_type) + except Exception as e: + log.error(f'[red]Error while loading {model_type} model: {context.model_paths[model_type]}[/red]') + log.error(f'[red]Error: {e}[/red]') + log.error(f'[red]Consider removing the model from the model folder.[red]') + + +def unload_all(context: Context): + for model_type in KNOWN_MODEL_TYPES: + unload_model(context, model_type) + +def resolve_model_to_use(model_name:str=None, model_type:str=None): + model_extensions = MODEL_EXTENSIONS.get(model_type, []) + default_models = DEFAULT_MODELS.get(model_type, []) + config = app.getConfig() + + model_dirs = [os.path.join(app.MODELS_DIR, model_type), app.SD_DIR] + if not model_name: # When None try user configured model. + # config = getConfig() + if 'model' in config and model_type in config['model']: + model_name = config['model'][model_type] + + if model_name: + # Check models directory + models_dir_path = os.path.join(app.MODELS_DIR, model_type, model_name) + for model_extension in model_extensions: + if os.path.exists(models_dir_path + model_extension): + return models_dir_path + model_extension + if os.path.exists(model_name + model_extension): + return os.path.abspath(model_name + model_extension) + + # Default locations + if model_name in default_models: + default_model_path = os.path.join(app.SD_DIR, model_name) + for model_extension in model_extensions: + if os.path.exists(default_model_path + model_extension): + return default_model_path + model_extension + + # Can't find requested model, check the default paths. + for default_model in default_models: + for model_dir in model_dirs: + default_model_path = os.path.join(model_dir, default_model) + for model_extension in model_extensions: + if os.path.exists(default_model_path + model_extension): + if model_name is not None: + log.warn(f'Could not find the configured custom model {model_name}{model_extension}. Using the default one: {default_model_path}{model_extension}') + return default_model_path + model_extension + + return None + +def reload_models_if_necessary(context: Context, task_data: TaskData): + model_paths_in_req = { + 'stable-diffusion': task_data.use_stable_diffusion_model, + 'vae': task_data.use_vae_model, + 'hypernetwork': task_data.use_hypernetwork_model, + 'gfpgan': task_data.use_face_correction, + 'realesrgan': task_data.use_upscale, + } + models_to_reload = {model_type: path for model_type, path in model_paths_in_req.items() if context.model_paths.get(model_type) != path} + + if set_vram_optimizations(context): # reload SD + models_to_reload['stable-diffusion'] = model_paths_in_req['stable-diffusion'] + + for model_type, model_path_in_req in models_to_reload.items(): + context.model_paths[model_type] = model_path_in_req + + action_fn = unload_model if context.model_paths[model_type] is None else load_model + action_fn(context, model_type, scan_model=False) # we've scanned them already + +def resolve_model_paths(task_data: TaskData): + task_data.use_stable_diffusion_model = resolve_model_to_use(task_data.use_stable_diffusion_model, model_type='stable-diffusion') + task_data.use_vae_model = resolve_model_to_use(task_data.use_vae_model, model_type='vae') + task_data.use_hypernetwork_model = resolve_model_to_use(task_data.use_hypernetwork_model, model_type='hypernetwork') + + if task_data.use_face_correction: task_data.use_face_correction = resolve_model_to_use(task_data.use_face_correction, 'gfpgan') + if task_data.use_upscale: task_data.use_upscale = resolve_model_to_use(task_data.use_upscale, 'realesrgan') + +def set_vram_optimizations(context: Context): + config = app.getConfig() + vram_usage_level = config.get('vram_usage_level', 'balanced') + + if vram_usage_level != context.vram_usage_level: + context.vram_usage_level = vram_usage_level + return True + + return False + +def make_model_folders(): + for model_type in KNOWN_MODEL_TYPES: + model_dir_path = os.path.join(app.MODELS_DIR, model_type) + + os.makedirs(model_dir_path, exist_ok=True) + + help_file_name = f'Place your {model_type} model files here.txt' + help_file_contents = f'Supported extensions: {" or ".join(MODEL_EXTENSIONS.get(model_type))}' + + with open(os.path.join(model_dir_path, help_file_name), 'w', encoding='utf-8') as f: + f.write(help_file_contents) + +def is_malicious_model(file_path): + try: + scan_result = scan_model(file_path) + if scan_result.issues_count > 0 or scan_result.infected_files > 0: + log.warn(":warning: [bold red]Scan %s: %d scanned, %d issue, %d infected.[/bold red]" % (file_path, scan_result.scanned_files, scan_result.issues_count, scan_result.infected_files)) + return True + else: + log.debug("Scan %s: [green]%d scanned, %d issue, %d infected.[/green]" % (file_path, scan_result.scanned_files, scan_result.issues_count, scan_result.infected_files)) + return False + except Exception as e: + log.error(f'error while scanning: {file_path}, error: {e}') + return False + +def getModels(): + models = { + 'active': { + 'stable-diffusion': 'sd-v1-4', + 'vae': '', + 'hypernetwork': '', + }, + 'options': { + 'stable-diffusion': ['sd-v1-4'], + 'vae': [], + 'hypernetwork': [], + }, + } + + models_scanned = 0 + + class MaliciousModelException(Exception): + "Raised when picklescan reports a problem with a model" + pass + + def scan_directory(directory, suffixes): + nonlocal models_scanned + tree = [] + for entry in os.scandir(directory): + if entry.is_file(): + matching_suffix = list(filter(lambda s: entry.name.endswith(s), suffixes)) + if len(matching_suffix) == 0: continue + matching_suffix = matching_suffix[0] + + mtime = entry.stat().st_mtime + mod_time = known_models[entry.path] if entry.path in known_models else -1 + if mod_time != mtime: + models_scanned += 1 + if is_malicious_model(entry.path): + raise MaliciousModelException(entry.path) + known_models[entry.path] = mtime + tree.append(entry.name[:-len(matching_suffix)]) + elif entry.is_dir(): + scan=scan_directory(entry.path, suffixes) + if len(scan) != 0: + tree.append( (entry.name, scan ) ) + return tree + + def listModels(model_type): + nonlocal models_scanned + + model_extensions = MODEL_EXTENSIONS.get(model_type, []) + models_dir = os.path.join(app.MODELS_DIR, model_type) + if not os.path.exists(models_dir): + os.makedirs(models_dir) + + try: + models['options'][model_type] = scan_directory(models_dir, model_extensions) + except MaliciousModelException as e: + models['scan-error'] = e + + # custom models + listModels(model_type='stable-diffusion') + listModels(model_type='vae') + listModels(model_type='hypernetwork') + + if models_scanned > 0: log.info(f'[green]Scanned {models_scanned} models. Nothing infected[/]') + + # legacy + custom_weight_path = os.path.join(app.SD_DIR, 'custom-model.ckpt') + if os.path.exists(custom_weight_path): + models['options']['stable-diffusion'].append('custom-model') + + return models diff --git a/ui/easydiffusion/renderer.py b/ui/easydiffusion/renderer.py new file mode 100644 index 00000000..672614ba --- /dev/null +++ b/ui/easydiffusion/renderer.py @@ -0,0 +1,137 @@ +import queue +import time +import json +import pprint + +from easydiffusion import device_manager +from easydiffusion.types import TaskData, Response, Image as ResponseImage, UserInitiatedStop, GenerateImageRequest +from easydiffusion.utils import get_printable_request, save_images_to_disk, log + +from sdkit import Context +from sdkit.generate import generate_images +from sdkit.filter import apply_filters +from sdkit.utils import img_to_buffer, img_to_base64_str, latent_samples_to_images, gc + +context = Context() # thread-local +''' +runtime data (bound locally to this thread), for e.g. device, references to loaded models, optimization flags etc +''' + +def init(device): + ''' + Initializes the fields that will be bound to this runtime's context, and sets the current torch device + ''' + context.stop_processing = False + context.temp_images = {} + context.partial_x_samples = None + + device_manager.device_init(context, device) + +def make_images(req: GenerateImageRequest, task_data: TaskData, data_queue: queue.Queue, task_temp_images: list, step_callback): + context.stop_processing = False + print_task_info(req, task_data) + + images, seeds = make_images_internal(req, task_data, data_queue, task_temp_images, step_callback) + + res = Response(req, task_data, images=construct_response(images, seeds, task_data, base_seed=req.seed)) + res = res.json() + data_queue.put(json.dumps(res)) + log.info('Task completed') + + return res + +def print_task_info(req: GenerateImageRequest, task_data: TaskData): + req_str = pprint.pformat(get_printable_request(req)).replace("[","\[") + task_str = pprint.pformat(task_data.dict()).replace("[","\[") + log.info(f'request: {req_str}') + log.info(f'task data: {task_str}') + +def make_images_internal(req: GenerateImageRequest, task_data: TaskData, data_queue: queue.Queue, task_temp_images: list, step_callback): + images, user_stopped = generate_images_internal(req, task_data, data_queue, task_temp_images, step_callback, task_data.stream_image_progress) + filtered_images = filter_images(task_data, images, user_stopped) + + if task_data.save_to_disk_path is not None: + save_images_to_disk(images, filtered_images, req, task_data) + + seeds = [*range(req.seed, req.seed + len(images))] + if task_data.show_only_filtered_image or filtered_images is images: + return filtered_images, seeds + else: + return images + filtered_images, seeds + seeds + +def generate_images_internal(req: GenerateImageRequest, task_data: TaskData, data_queue: queue.Queue, task_temp_images: list, step_callback, stream_image_progress: bool): + context.temp_images.clear() + + callback = make_step_callback(req, task_data, data_queue, task_temp_images, step_callback, stream_image_progress) + + try: + if req.init_image is not None: req.sampler_name = 'ddim' + + images = generate_images(context, callback=callback, **req.dict()) + user_stopped = False + except UserInitiatedStop: + images = [] + user_stopped = True + if context.partial_x_samples is not None: + images = latent_samples_to_images(context, context.partial_x_samples) + finally: + if hasattr(context, 'partial_x_samples') and context.partial_x_samples is not None: + del context.partial_x_samples + context.partial_x_samples = None + + return images, user_stopped + +def filter_images(task_data: TaskData, images: list, user_stopped): + if user_stopped or (task_data.use_face_correction is None and task_data.use_upscale is None): + return images + + filters_to_apply = [] + if task_data.use_face_correction and 'gfpgan' in task_data.use_face_correction.lower(): filters_to_apply.append('gfpgan') + if task_data.use_upscale and 'realesrgan' in task_data.use_upscale.lower(): filters_to_apply.append('realesrgan') + + return apply_filters(context, filters_to_apply, images, scale=task_data.upscale_amount) + +def construct_response(images: list, seeds: list, task_data: TaskData, base_seed: int): + return [ + ResponseImage( + data=img_to_base64_str(img, task_data.output_format, task_data.output_quality), + seed=seed, + ) for img, seed in zip(images, seeds) + ] + +def make_step_callback(req: GenerateImageRequest, task_data: TaskData, data_queue: queue.Queue, task_temp_images: list, step_callback, stream_image_progress: bool): + n_steps = req.num_inference_steps if req.init_image is None else int(req.num_inference_steps * req.prompt_strength) + last_callback_time = -1 + + def update_temp_img(x_samples, task_temp_images: list): + partial_images = [] + images = latent_samples_to_images(context, x_samples) + for i, img in enumerate(images): + buf = img_to_buffer(img, output_format='JPEG') + + context.temp_images[f"{task_data.request_id}/{i}"] = buf + task_temp_images[i] = buf + partial_images.append({'path': f"/image/tmp/{task_data.request_id}/{i}"}) + del images + return partial_images + + def on_image_step(x_samples, i): + nonlocal last_callback_time + + context.partial_x_samples = x_samples + step_time = time.time() - last_callback_time if last_callback_time != -1 else -1 + last_callback_time = time.time() + + progress = {"step": i, "step_time": step_time, "total_steps": n_steps} + + if stream_image_progress and i % 5 == 0: + progress['output'] = update_temp_img(x_samples, task_temp_images) + + data_queue.put(json.dumps(progress)) + + step_callback() + + if context.stop_processing: + raise UserInitiatedStop("User requested that we stop processing") + + return on_image_step diff --git a/ui/easydiffusion/server.py b/ui/easydiffusion/server.py new file mode 100644 index 00000000..7f32cffd --- /dev/null +++ b/ui/easydiffusion/server.py @@ -0,0 +1,241 @@ +"""server.py: FastAPI SD-UI Web Host. +Notes: + async endpoints always run on the main thread. Without they run on the thread pool. +""" +import os +import traceback +import datetime +from typing import List, Union + +from fastapi import FastAPI, HTTPException +from fastapi.staticfiles import StaticFiles +from starlette.responses import FileResponse, JSONResponse, StreamingResponse +from pydantic import BaseModel + +from easydiffusion import app, model_manager, task_manager +from easydiffusion.types import TaskData, GenerateImageRequest, MergeRequest +from easydiffusion.utils import log + +log.info(f'started in {app.SD_DIR}') +log.info(f'started at {datetime.datetime.now():%x %X}') + +server_api = FastAPI() + +NOCACHE_HEADERS={"Cache-Control": "no-cache, no-store, must-revalidate", "Pragma": "no-cache", "Expires": "0"} + +class NoCacheStaticFiles(StaticFiles): + def is_not_modified(self, response_headers, request_headers) -> bool: + if 'content-type' in response_headers and ('javascript' in response_headers['content-type'] or 'css' in response_headers['content-type']): + response_headers.update(NOCACHE_HEADERS) + return False + + return super().is_not_modified(response_headers, request_headers) + +class SetAppConfigRequest(BaseModel): + update_branch: str = None + render_devices: Union[List[str], List[int], str, int] = None + model_vae: str = None + ui_open_browser_on_start: bool = None + listen_to_network: bool = None + listen_port: int = None + +def init(): + server_api.mount('/media', NoCacheStaticFiles(directory=os.path.join(app.SD_UI_DIR, 'media')), name="media") + + for plugins_dir, dir_prefix in app.UI_PLUGINS_SOURCES: + server_api.mount(f'/plugins/{dir_prefix}', NoCacheStaticFiles(directory=plugins_dir), name=f"plugins-{dir_prefix}") + + @server_api.post('/app_config') + async def set_app_config(req : SetAppConfigRequest): + return set_app_config_internal(req) + + @server_api.get('/get/{key:path}') + def read_web_data(key:str=None): + return read_web_data_internal(key) + + @server_api.get('/ping') # Get server and optionally session status. + def ping(session_id:str=None): + return ping_internal(session_id) + + @server_api.post('/render') + def render(req: dict): + return render_internal(req) + + @server_api.post('/model/merge') + def model_merge(req: dict): + print(req) + return model_merge_internal(req) + + @server_api.get('/image/stream/{task_id:int}') + def stream(task_id:int): + return stream_internal(task_id) + + @server_api.get('/image/stop') + def stop(task: int): + return stop_internal(task) + + @server_api.get('/image/tmp/{task_id:int}/{img_id:int}') + def get_image(task_id: int, img_id: int): + return get_image_internal(task_id, img_id) + + @server_api.get('/') + def read_root(): + return FileResponse(os.path.join(app.SD_UI_DIR, 'index.html'), headers=NOCACHE_HEADERS) + + @server_api.on_event("shutdown") + def shutdown_event(): # Signal render thread to close on shutdown + task_manager.current_state_error = SystemExit('Application shutting down.') + +# API implementations +def set_app_config_internal(req : SetAppConfigRequest): + config = app.getConfig() + if req.update_branch is not None: + config['update_branch'] = req.update_branch + if req.render_devices is not None: + update_render_devices_in_config(config, req.render_devices) + if req.ui_open_browser_on_start is not None: + if 'ui' not in config: + config['ui'] = {} + config['ui']['open_browser_on_start'] = req.ui_open_browser_on_start + if req.listen_to_network is not None: + if 'net' not in config: + config['net'] = {} + config['net']['listen_to_network'] = bool(req.listen_to_network) + if req.listen_port is not None: + if 'net' not in config: + config['net'] = {} + config['net']['listen_port'] = int(req.listen_port) + try: + app.setConfig(config) + + if req.render_devices: + app.update_render_threads() + + return JSONResponse({'status': 'OK'}, headers=NOCACHE_HEADERS) + except Exception as e: + log.error(traceback.format_exc()) + raise HTTPException(status_code=500, detail=str(e)) + +def update_render_devices_in_config(config, render_devices): + if render_devices not in ('cpu', 'auto') and not render_devices.startswith('cuda:'): + raise HTTPException(status_code=400, detail=f'Invalid render device requested: {render_devices}') + + if render_devices.startswith('cuda:'): + render_devices = render_devices.split(',') + + config['render_devices'] = render_devices + +def read_web_data_internal(key:str=None): + if not key: # /get without parameters, stable-diffusion easter egg. + raise HTTPException(status_code=418, detail="StableDiffusion is drawing a teapot!") # HTTP418 I'm a teapot + elif key == 'app_config': + return JSONResponse(app.getConfig(), headers=NOCACHE_HEADERS) + elif key == 'system_info': + config = app.getConfig() + system_info = { + 'devices': task_manager.get_devices(), + 'hosts': app.getIPConfig(), + 'default_output_dir': os.path.join(os.path.expanduser("~"), app.OUTPUT_DIRNAME), + } + system_info['devices']['config'] = config.get('render_devices', "auto") + return JSONResponse(system_info, headers=NOCACHE_HEADERS) + elif key == 'models': + return JSONResponse(model_manager.getModels(), headers=NOCACHE_HEADERS) + elif key == 'modifiers': return FileResponse(os.path.join(app.SD_UI_DIR, 'modifiers.json'), headers=NOCACHE_HEADERS) + elif key == 'ui_plugins': return JSONResponse(app.getUIPlugins(), headers=NOCACHE_HEADERS) + else: + raise HTTPException(status_code=404, detail=f'Request for unknown {key}') # HTTP404 Not Found + +def ping_internal(session_id:str=None): + if task_manager.is_alive() <= 0: # Check that render threads are alive. + if task_manager.current_state_error: raise HTTPException(status_code=500, detail=str(task_manager.current_state_error)) + raise HTTPException(status_code=500, detail='Render thread is dead.') + if task_manager.current_state_error and not isinstance(task_manager.current_state_error, StopAsyncIteration): raise HTTPException(status_code=500, detail=str(task_manager.current_state_error)) + # Alive + response = {'status': str(task_manager.current_state)} + if session_id: + session = task_manager.get_cached_session(session_id, update_ttl=True) + response['tasks'] = {id(t): t.status for t in session.tasks} + response['devices'] = task_manager.get_devices() + return JSONResponse(response, headers=NOCACHE_HEADERS) + +def render_internal(req: dict): + try: + # separate out the request data into rendering and task-specific data + render_req: GenerateImageRequest = GenerateImageRequest.parse_obj(req) + task_data: TaskData = TaskData.parse_obj(req) + + render_req.init_image_mask = req.get('mask') # hack: will rename this in the HTTP API in a future revision + + app.save_to_config(task_data.use_stable_diffusion_model, task_data.use_vae_model, task_data.use_hypernetwork_model, task_data.vram_usage_level) + + # enqueue the task + new_task = task_manager.render(render_req, task_data) + response = { + 'status': str(task_manager.current_state), + 'queue': len(task_manager.tasks_queue), + 'stream': f'/image/stream/{id(new_task)}', + 'task': id(new_task) + } + return JSONResponse(response, headers=NOCACHE_HEADERS) + except ChildProcessError as e: # Render thread is dead + raise HTTPException(status_code=500, detail=f'Rendering thread has died.') # HTTP500 Internal Server Error + except ConnectionRefusedError as e: # Unstarted task pending limit reached, deny queueing too many. + raise HTTPException(status_code=503, detail=str(e)) # HTTP503 Service Unavailable + except Exception as e: + log.error(traceback.format_exc()) + raise HTTPException(status_code=500, detail=str(e)) + +def model_merge_internal(req: dict): + try: + from sdkit.train import merge_models + from easydiffusion.utils.save_utils import filename_regex + mergeReq: MergeRequest = MergeRequest.parse_obj(req) + + merge_models(model_manager.resolve_model_to_use(mergeReq.model0,'stable-diffusion'), + model_manager.resolve_model_to_use(mergeReq.model1,'stable-diffusion'), + mergeReq.ratio, + os.path.join(app.MODELS_DIR, 'stable-diffusion', filename_regex.sub('_', mergeReq.out_path)), + mergeReq.use_fp16 + ) + return JSONResponse({'status':'OK'}, headers=NOCACHE_HEADERS) + except Exception as e: + log.error(traceback.format_exc()) + raise HTTPException(status_code=500, detail=str(e)) + +def stream_internal(task_id:int): + #TODO Move to WebSockets ?? + task = task_manager.get_cached_task(task_id, update_ttl=True) + if not task: raise HTTPException(status_code=404, detail=f'Request {task_id} not found.') # HTTP404 NotFound + #if (id(task) != task_id): raise HTTPException(status_code=409, detail=f'Wrong task id received. Expected:{id(task)}, Received:{task_id}') # HTTP409 Conflict + if task.buffer_queue.empty() and not task.lock.locked(): + if task.response: + #log.info(f'Session {session_id} sending cached response') + return JSONResponse(task.response, headers=NOCACHE_HEADERS) + raise HTTPException(status_code=425, detail='Too Early, task not started yet.') # HTTP425 Too Early + #log.info(f'Session {session_id} opened live render stream {id(task.buffer_queue)}') + return StreamingResponse(task.read_buffer_generator(), media_type='application/json') + +def stop_internal(task: int): + if not task: + if task_manager.current_state == task_manager.ServerStates.Online or task_manager.current_state == task_manager.ServerStates.Unavailable: + raise HTTPException(status_code=409, detail='Not currently running any tasks.') # HTTP409 Conflict + task_manager.current_state_error = StopAsyncIteration('') + return {'OK'} + task_id = task + task = task_manager.get_cached_task(task_id, update_ttl=False) + if not task: raise HTTPException(status_code=404, detail=f'Task {task_id} was not found.') # HTTP404 Not Found + if isinstance(task.error, StopAsyncIteration): raise HTTPException(status_code=409, detail=f'Task {task_id} is already stopped.') # HTTP409 Conflict + task.error = StopAsyncIteration(f'Task {task_id} stop requested.') + return {'OK'} + +def get_image_internal(task_id: int, img_id: int): + task = task_manager.get_cached_task(task_id, update_ttl=True) + if not task: raise HTTPException(status_code=410, detail=f'Task {task_id} could not be found.') # HTTP404 NotFound + if not task.temp_images[img_id]: raise HTTPException(status_code=425, detail='Too Early, task data is not available yet.') # HTTP425 Too Early + try: + img_data = task.temp_images[img_id] + img_data.seek(0) + return StreamingResponse(img_data, media_type='image/jpeg') + except KeyError as e: + raise HTTPException(status_code=500, detail=str(e)) diff --git a/ui/sd_internal/task_manager.py b/ui/easydiffusion/task_manager.py similarity index 65% rename from ui/sd_internal/task_manager.py rename to ui/easydiffusion/task_manager.py index a8c5a4b2..213f5a96 100644 --- a/ui/sd_internal/task_manager.py +++ b/ui/easydiffusion/task_manager.py @@ -11,12 +11,15 @@ TASK_TTL = 15 * 60 # seconds, Discard last session's task timeout import torch import queue, threading, time, weakref -from typing import Any, Generator, Hashable, Optional, Union +from typing import Any, Hashable -from pydantic import BaseModel -from sd_internal import Request, Response, runtime, device_manager +from easydiffusion import device_manager +from easydiffusion.types import TaskData, GenerateImageRequest +from easydiffusion.utils import log -THREAD_NAME_PREFIX = 'Runtime-Render/' +from sdkit.utils import gc + +THREAD_NAME_PREFIX = '' ERR_LOCK_FAILED = ' failed to acquire lock within timeout.' LOCK_TIMEOUT = 15 # Maximum locking time in seconds before failing a task. # It's better to get an exception than a deadlock... ALWAYS use timeout in critical paths. @@ -36,12 +39,13 @@ class ServerStates: class Unavailable(Symbol): pass class RenderTask(): # Task with output queue and completion lock. - def __init__(self, req: Request): - req.request_id = id(self) - self.request: Request = req # Initial Request + def __init__(self, req: GenerateImageRequest, task_data: TaskData): + task_data.request_id = id(self) + self.render_request: GenerateImageRequest = req # Initial Request + self.task_data: TaskData = task_data self.response: Any = None # Copy of the last reponse self.render_device = None # Select the task affinity. (Not used to change active devices). - self.temp_images:list = [None] * req.num_outputs * (1 if req.show_only_filtered_image else 2) + self.temp_images:list = [None] * req.num_outputs * (1 if task_data.show_only_filtered_image else 2) self.error: Exception = None self.lock: threading.Lock = threading.Lock() # Locks at task start and unlocks when task is completed self.buffer_queue: queue.Queue = queue.Queue() # Queue of JSON string segments @@ -69,54 +73,6 @@ class RenderTask(): # Task with output queue and completion lock. def is_pending(self): return bool(not self.response and not self.error) -# defaults from https://huggingface.co/blog/stable_diffusion -class ImageRequest(BaseModel): - session_id: str = "session" - prompt: str = "" - negative_prompt: str = "" - init_image: str = None # base64 - mask: str = None # base64 - num_outputs: int = 1 - num_inference_steps: int = 50 - guidance_scale: float = 7.5 - width: int = 512 - height: int = 512 - seed: int = 42 - prompt_strength: float = 0.8 - sampler: str = None # "ddim", "plms", "heun", "euler", "euler_a", "dpm2", "dpm2_a", "lms" - # allow_nsfw: bool = False - save_to_disk_path: str = None - turbo: bool = True - use_cpu: bool = False ##TODO Remove after UI and plugins transition. - render_device: str = None # Select the task affinity. (Not used to change active devices). - use_full_precision: bool = False - use_face_correction: str = None # or "GFPGANv1.3" - use_upscale: str = None # or "RealESRGAN_x4plus" or "RealESRGAN_x4plus_anime_6B" - use_stable_diffusion_model: str = "sd-v1-4" - use_vae_model: str = None - use_hypernetwork_model: str = None - hypernetwork_strength: float = None - show_only_filtered_image: bool = False - output_format: str = "jpeg" # or "png" - output_quality: int = 75 - - stream_progress_updates: bool = False - stream_image_progress: bool = False - -class FilterRequest(BaseModel): - session_id: str = "session" - model: str = None - name: str = "" - init_image: str = None # base64 - width: int = 512 - height: int = 512 - save_to_disk_path: str = None - turbo: bool = True - render_device: str = None - use_full_precision: bool = False - output_format: str = "jpeg" # or "png" - output_quality: int = 75 - # Temporary cache to allow to query tasks results for a short time after they are completed. class DataCache(): def __init__(self): @@ -139,11 +95,11 @@ class DataCache(): for key in to_delete: (_, val) = self._base[key] if isinstance(val, RenderTask): - print(f'RenderTask {key} expired. Data removed.') + log.debug(f'RenderTask {key} expired. Data removed.') elif isinstance(val, SessionState): - print(f'Session {key} expired. Data removed.') + log.debug(f'Session {key} expired. Data removed.') else: - print(f'Key {key} expired. Data removed.') + log.debug(f'Key {key} expired. Data removed.') del self._base[key] finally: self._lock.release() @@ -177,8 +133,7 @@ class DataCache(): self._get_ttl_time(ttl), value ) except Exception as e: - print(str(e)) - print(traceback.format_exc()) + log.error(traceback.format_exc()) return False else: return True @@ -189,7 +144,7 @@ class DataCache(): try: ttl, value = self._base.get(key, (None, None)) if ttl is not None and self._is_expired(ttl): - print(f'Session {key} expired. Discarding data.') + log.debug(f'Session {key} expired. Discarding data.') del self._base[key] return None return value @@ -200,15 +155,9 @@ manager_lock = threading.RLock() render_threads = [] current_state = ServerStates.Init current_state_error:Exception = None -current_model_path = None -current_vae_path = None -current_hypernetwork_path = None tasks_queue = [] session_cache = DataCache() task_cache = DataCache() -default_model_to_load = None -default_vae_to_load = None -default_hypernetwork_to_load = None weak_thread_data = weakref.WeakKeyDictionary() idle_event: threading.Event = threading.Event() @@ -236,40 +185,10 @@ class SessionState(): self._tasks_ids.pop(0) return True -def preload_model(ckpt_file_path=None, vae_file_path=None, hypernetwork_file_path=None): - global current_state, current_state_error, current_model_path, current_vae_path, current_hypernetwork_path - if ckpt_file_path == None: - ckpt_file_path = default_model_to_load - if vae_file_path == None: - vae_file_path = default_vae_to_load - if hypernetwork_file_path == None: - hypernetwork_file_path = default_hypernetwork_to_load - if ckpt_file_path == current_model_path and vae_file_path == current_vae_path: - return - current_state = ServerStates.LoadingModel - try: - from . import runtime - runtime.thread_data.hypernetwork_file = hypernetwork_file_path - runtime.thread_data.ckpt_file = ckpt_file_path - runtime.thread_data.vae_file = vae_file_path - runtime.load_model_ckpt() - runtime.load_hypernetwork() - current_model_path = ckpt_file_path - current_vae_path = vae_file_path - current_hypernetwork_path = hypernetwork_file_path - current_state_error = None - current_state = ServerStates.Online - except Exception as e: - current_model_path = None - current_vae_path = None - current_state_error = e - current_state = ServerStates.Unavailable - print(traceback.format_exc()) - def thread_get_next_task(): - from . import runtime + from easydiffusion import renderer if not manager_lock.acquire(blocking=True, timeout=LOCK_TIMEOUT): - print('Render thread on device', runtime.thread_data.device, 'failed to acquire manager lock.') + log.warn(f'Render thread on device: {renderer.context.device} failed to acquire manager lock.') return None if len(tasks_queue) <= 0: manager_lock.release() @@ -277,7 +196,7 @@ def thread_get_next_task(): task = None try: # Select a render task. for queued_task in tasks_queue: - if queued_task.render_device and runtime.thread_data.device != queued_task.render_device: + if queued_task.render_device and renderer.context.device != queued_task.render_device: # Is asking for a specific render device. if is_alive(queued_task.render_device) > 0: continue # requested device alive, skip current one. @@ -286,7 +205,7 @@ def thread_get_next_task(): queued_task.error = Exception(queued_task.render_device + ' is not currently active.') task = queued_task break - if not queued_task.render_device and runtime.thread_data.device == 'cpu' and is_alive() > 1: + if not queued_task.render_device and renderer.context.device == 'cpu' and is_alive() > 1: # not asking for any specific devices, cpu want to grab task but other render devices are alive. continue # Skip Tasks, don't run on CPU unless there is nothing else or user asked for it. task = queued_task @@ -298,31 +217,36 @@ def thread_get_next_task(): manager_lock.release() def thread_render(device): - global current_state, current_state_error, current_model_path, current_vae_path, current_hypernetwork_path - from . import runtime + global current_state, current_state_error + + from easydiffusion import renderer, model_manager try: - runtime.thread_init(device) - except Exception as e: - print(traceback.format_exc()) + renderer.init(device) + weak_thread_data[threading.current_thread()] = { - 'error': e + 'device': renderer.context.device, + 'device_name': renderer.context.device_name, + 'alive': True + } + + current_state = ServerStates.LoadingModel + model_manager.load_default_models(renderer.context) + + current_state = ServerStates.Online + except Exception as e: + log.error(traceback.format_exc()) + weak_thread_data[threading.current_thread()] = { + 'error': e, + 'alive': False } return - weak_thread_data[threading.current_thread()] = { - 'device': runtime.thread_data.device, - 'device_name': runtime.thread_data.device_name, - 'alive': True - } - if runtime.thread_data.device != 'cpu' or is_alive() == 1: - preload_model() - current_state = ServerStates.Online + while True: session_cache.clean() task_cache.clean() if not weak_thread_data[threading.current_thread()]['alive']: - print(f'Shutting down thread for device {runtime.thread_data.device}') - runtime.unload_models() - runtime.unload_filters() + log.info(f'Shutting down thread for device {renderer.context.device}') + model_manager.unload_all(renderer.context) return if isinstance(current_state_error, SystemExit): current_state = ServerStates.Unavailable @@ -333,7 +257,7 @@ def thread_render(device): idle_event.wait(timeout=1) continue if task.error is not None: - print(task.error) + log.error(task.error) task.response = {"status": 'failed', "detail": str(task.error)} task.buffer_queue.put(json.dumps(task.response)) continue @@ -342,51 +266,44 @@ def thread_render(device): task.response = {"status": 'failed', "detail": str(task.error)} task.buffer_queue.put(json.dumps(task.response)) continue - print(f'Session {task.request.session_id} starting task {id(task)} on {runtime.thread_data.device_name}') + log.info(f'Session {task.task_data.session_id} starting task {id(task)} on {renderer.context.device_name}') if not task.lock.acquire(blocking=False): raise Exception('Got locked task from queue.') try: - if runtime.is_hypernetwork_reload_necessary(task.request): - runtime.reload_hypernetwork() - current_hypernetwork_path = task.request.use_hypernetwork_model - - if runtime.is_model_reload_necessary(task.request): - current_state = ServerStates.LoadingModel - runtime.reload_model() - current_model_path = task.request.use_stable_diffusion_model - current_vae_path = task.request.use_vae_model - def step_callback(): global current_state_error if isinstance(current_state_error, SystemExit) or isinstance(current_state_error, StopAsyncIteration) or isinstance(task.error, StopAsyncIteration): - runtime.thread_data.stop_processing = True + renderer.context.stop_processing = True if isinstance(current_state_error, StopAsyncIteration): task.error = current_state_error current_state_error = None - print(f'Session {task.request.session_id} sent cancel signal for task {id(task)}') + log.info(f'Session {task.task_data.session_id} sent cancel signal for task {id(task)}') + + current_state = ServerStates.LoadingModel + model_manager.resolve_model_paths(task.task_data) + model_manager.reload_models_if_necessary(renderer.context, task.task_data) current_state = ServerStates.Rendering - task.response = runtime.mk_img(task.request, task.buffer_queue, task.temp_images, step_callback) + task.response = renderer.make_images(task.render_request, task.task_data, task.buffer_queue, task.temp_images, step_callback) # Before looping back to the generator, mark cache as still alive. task_cache.keep(id(task), TASK_TTL) - session_cache.keep(task.request.session_id, TASK_TTL) + session_cache.keep(task.task_data.session_id, TASK_TTL) except Exception as e: - task.error = e + task.error = str(e) task.response = {"status": 'failed', "detail": str(task.error)} task.buffer_queue.put(json.dumps(task.response)) - print(traceback.format_exc()) - continue + log.error(traceback.format_exc()) finally: - # Task completed + gc(renderer.context) task.lock.release() task_cache.keep(id(task), TASK_TTL) - session_cache.keep(task.request.session_id, TASK_TTL) + session_cache.keep(task.task_data.session_id, TASK_TTL) if isinstance(task.error, StopAsyncIteration): - print(f'Session {task.request.session_id} task {id(task)} cancelled!') + log.info(f'Session {task.task_data.session_id} task {id(task)} cancelled!') elif task.error is not None: - print(f'Session {task.request.session_id} task {id(task)} failed!') + log.info(f'Session {task.task_data.session_id} task {id(task)} failed!') else: - print(f'Session {task.request.session_id} task {id(task)} completed by {runtime.thread_data.device_name}.') + log.info(f'Session {task.task_data.session_id} task {id(task)} completed by {renderer.context.device_name}.') current_state = ServerStates.Online def get_cached_task(task_id:str, update_ttl:bool=False): @@ -423,6 +340,7 @@ def get_devices(): 'name': torch.cuda.get_device_name(device), 'mem_free': mem_free, 'mem_total': mem_total, + 'max_vram_usage_level': device_manager.get_max_vram_usage_level(device), } # list the compatible devices @@ -472,7 +390,7 @@ def is_alive(device=None): def start_render_thread(device): if not manager_lock.acquire(blocking=True, timeout=LOCK_TIMEOUT): raise Exception('start_render_thread' + ERR_LOCK_FAILED) - print('Start new Rendering Thread on device', device) + log.info(f'Start new Rendering Thread on device: {device}') try: rthread = threading.Thread(target=thread_render, kwargs={'device': device}) rthread.daemon = True @@ -484,7 +402,7 @@ def start_render_thread(device): timeout = DEVICE_START_TIMEOUT while not rthread.is_alive() or not rthread in weak_thread_data or not 'device' in weak_thread_data[rthread]: if rthread in weak_thread_data and 'error' in weak_thread_data[rthread]: - print(rthread, device, 'error:', weak_thread_data[rthread]['error']) + log.error(f"{rthread}, {device}, error: {weak_thread_data[rthread]['error']}") return False if timeout <= 0: return False @@ -496,11 +414,11 @@ def stop_render_thread(device): try: device_manager.validate_device_id(device, log_prefix='stop_render_thread') except: - print(traceback.format_exc()) + log.error(traceback.format_exc()) return False if not manager_lock.acquire(blocking=True, timeout=LOCK_TIMEOUT): raise Exception('stop_render_thread' + ERR_LOCK_FAILED) - print('Stopping Rendering Thread on device', device) + log.info(f'Stopping Rendering Thread on device: {device}') try: thread_to_remove = None @@ -523,79 +441,44 @@ def stop_render_thread(device): def update_render_threads(render_devices, active_devices): devices_to_start, devices_to_stop = device_manager.get_device_delta(render_devices, active_devices) - print('devices_to_start', devices_to_start) - print('devices_to_stop', devices_to_stop) + log.debug(f'devices_to_start: {devices_to_start}') + log.debug(f'devices_to_stop: {devices_to_stop}') for device in devices_to_stop: if is_alive(device) <= 0: - print(device, 'is not alive') + log.debug(f'{device} is not alive') continue if not stop_render_thread(device): - print(device, 'could not stop render thread') + log.warn(f'{device} could not stop render thread') for device in devices_to_start: if is_alive(device) >= 1: - print(device, 'already registered.') + log.debug(f'{device} already registered.') continue if not start_render_thread(device): - print(device, 'failed to start.') + log.warn(f'{device} failed to start.') if is_alive() <= 0: # No running devices, probably invalid user config. raise EnvironmentError('ERROR: No active render devices! Please verify the "render_devices" value in config.json') - print('active devices', get_devices()['active']) + log.debug(f"active devices: {get_devices()['active']}") def shutdown_event(): # Signal render thread to close on shutdown global current_state_error current_state_error = SystemExit('Application shutting down.') -def render(req : ImageRequest): +def render(render_req: GenerateImageRequest, task_data: TaskData): current_thread_count = is_alive() if current_thread_count <= 0: # Render thread is dead raise ChildProcessError('Rendering thread has died.') # Alive, check if task in cache - session = get_cached_session(req.session_id, update_ttl=True) + session = get_cached_session(task_data.session_id, update_ttl=True) pending_tasks = list(filter(lambda t: t.is_pending, session.tasks)) if current_thread_count < len(pending_tasks): - raise ConnectionRefusedError(f'Session {req.session_id} already has {len(pending_tasks)} pending tasks out of {current_thread_count}.') + raise ConnectionRefusedError(f'Session {task_data.session_id} already has {len(pending_tasks)} pending tasks out of {current_thread_count}.') - from . import runtime - r = Request() - r.session_id = req.session_id - r.prompt = req.prompt - r.negative_prompt = req.negative_prompt - r.init_image = req.init_image - r.mask = req.mask - r.num_outputs = req.num_outputs - r.num_inference_steps = req.num_inference_steps - r.guidance_scale = req.guidance_scale - r.width = req.width - r.height = req.height - r.seed = req.seed - r.prompt_strength = req.prompt_strength - r.sampler = req.sampler - # r.allow_nsfw = req.allow_nsfw - r.turbo = req.turbo - r.use_full_precision = req.use_full_precision - r.save_to_disk_path = req.save_to_disk_path - r.use_upscale: str = req.use_upscale - r.use_face_correction = req.use_face_correction - r.use_stable_diffusion_model = req.use_stable_diffusion_model - r.use_vae_model = req.use_vae_model - r.use_hypernetwork_model = req.use_hypernetwork_model - r.hypernetwork_strength = req.hypernetwork_strength - r.show_only_filtered_image = req.show_only_filtered_image - r.output_format = req.output_format - r.output_quality = req.output_quality - - r.stream_progress_updates = True # the underlying implementation only supports streaming - r.stream_image_progress = req.stream_image_progress - - if not req.stream_progress_updates: - r.stream_image_progress = False - - new_task = RenderTask(r) + new_task = RenderTask(render_req, task_data) if session.put(new_task, TASK_TTL): # Use twice the normal timeout for adding user requests. # Tries to force session.put to fail before tasks_queue.put would. diff --git a/ui/easydiffusion/types.py b/ui/easydiffusion/types.py new file mode 100644 index 00000000..c70e125b --- /dev/null +++ b/ui/easydiffusion/types.py @@ -0,0 +1,95 @@ +from pydantic import BaseModel +from typing import Any + +class GenerateImageRequest(BaseModel): + prompt: str = "" + negative_prompt: str = "" + + seed: int = 42 + width: int = 512 + height: int = 512 + + num_outputs: int = 1 + num_inference_steps: int = 50 + guidance_scale: float = 7.5 + + init_image: Any = None + init_image_mask: Any = None + prompt_strength: float = 0.8 + preserve_init_image_color_profile = False + + sampler_name: str = None # "ddim", "plms", "heun", "euler", "euler_a", "dpm2", "dpm2_a", "lms" + hypernetwork_strength: float = 0 + +class TaskData(BaseModel): + request_id: str = None + session_id: str = "session" + save_to_disk_path: str = None + vram_usage_level: str = "balanced" # or "low" or "medium" + + use_face_correction: str = None # or "GFPGANv1.3" + use_upscale: str = None # or "RealESRGAN_x4plus" or "RealESRGAN_x4plus_anime_6B" + upscale_amount: int = 4 # or 2 + use_stable_diffusion_model: str = "sd-v1-4" + # use_stable_diffusion_config: str = "v1-inference" + use_vae_model: str = None + use_hypernetwork_model: str = None + + show_only_filtered_image: bool = False + output_format: str = "jpeg" # or "png" + output_quality: int = 75 + metadata_output_format: str = "txt" # or "json" + stream_image_progress: bool = False + +class MergeRequest(BaseModel): + model0: str = None + model1: str = None + ratio: float = None + out_path: str = "mix" + use_fp16 = True + +class Image: + data: str # base64 + seed: int + is_nsfw: bool + path_abs: str = None + + def __init__(self, data, seed): + self.data = data + self.seed = seed + + def json(self): + return { + "data": self.data, + "seed": self.seed, + "path_abs": self.path_abs, + } + +class Response: + render_request: GenerateImageRequest + task_data: TaskData + images: list + + def __init__(self, render_request: GenerateImageRequest, task_data: TaskData, images: list): + self.render_request = render_request + self.task_data = task_data + self.images = images + + def json(self): + del self.render_request.init_image + del self.render_request.init_image_mask + + res = { + "status": 'succeeded', + "render_request": self.render_request.dict(), + "task_data": self.task_data.dict(), + "output": [], + } + + for image in self.images: + res["output"].append(image.json()) + + return res + +class UserInitiatedStop(Exception): + pass diff --git a/ui/easydiffusion/utils/__init__.py b/ui/easydiffusion/utils/__init__.py new file mode 100644 index 00000000..8be070b4 --- /dev/null +++ b/ui/easydiffusion/utils/__init__.py @@ -0,0 +1,8 @@ +import logging + +log = logging.getLogger('easydiffusion') + +from .save_utils import ( + save_images_to_disk, + get_printable_request, +) \ No newline at end of file diff --git a/ui/easydiffusion/utils/save_utils.py b/ui/easydiffusion/utils/save_utils.py new file mode 100644 index 00000000..0b9e3bce --- /dev/null +++ b/ui/easydiffusion/utils/save_utils.py @@ -0,0 +1,88 @@ +import os +import time +import base64 +import re + +from easydiffusion.types import TaskData, GenerateImageRequest + +from sdkit.utils import save_images, save_dicts + +filename_regex = re.compile('[^a-zA-Z0-9._-]') + +# keep in sync with `ui/media/js/dnd.js` +TASK_TEXT_MAPPING = { + 'prompt': 'Prompt', + 'width': 'Width', + 'height': 'Height', + 'seed': 'Seed', + 'num_inference_steps': 'Steps', + 'guidance_scale': 'Guidance Scale', + 'prompt_strength': 'Prompt Strength', + 'use_face_correction': 'Use Face Correction', + 'use_upscale': 'Use Upscaling', + 'upscale_amount': 'Upscale By', + 'sampler_name': 'Sampler', + 'negative_prompt': 'Negative Prompt', + 'use_stable_diffusion_model': 'Stable Diffusion model', + 'use_hypernetwork_model': 'Hypernetwork model', + 'hypernetwork_strength': 'Hypernetwork Strength' +} + +def save_images_to_disk(images: list, filtered_images: list, req: GenerateImageRequest, task_data: TaskData): + now = time.time() + save_dir_path = os.path.join(task_data.save_to_disk_path, filename_regex.sub('_', task_data.session_id)) + metadata_entries = get_metadata_entries_for_request(req, task_data) + make_filename = make_filename_callback(req, now=now) + + if task_data.show_only_filtered_image or filtered_images is images: + save_images(filtered_images, save_dir_path, file_name=make_filename, output_format=task_data.output_format, output_quality=task_data.output_quality) + save_dicts(metadata_entries, save_dir_path, file_name=make_filename, output_format=task_data.metadata_output_format) + else: + make_filter_filename = make_filename_callback(req, now=now, suffix='filtered') + + save_images(images, save_dir_path, file_name=make_filename, output_format=task_data.output_format, output_quality=task_data.output_quality) + save_images(filtered_images, save_dir_path, file_name=make_filter_filename, output_format=task_data.output_format, output_quality=task_data.output_quality) + save_dicts(metadata_entries, save_dir_path, file_name=make_filter_filename, output_format=task_data.metadata_output_format) + +def get_metadata_entries_for_request(req: GenerateImageRequest, task_data: TaskData): + metadata = get_printable_request(req) + metadata.update({ + 'use_stable_diffusion_model': task_data.use_stable_diffusion_model, + 'use_vae_model': task_data.use_vae_model, + 'use_hypernetwork_model': task_data.use_hypernetwork_model, + 'use_face_correction': task_data.use_face_correction, + 'use_upscale': task_data.use_upscale, + }) + if metadata['use_upscale'] is not None: + metadata['upscale_amount'] = task_data.upscale_amount + + # if text, format it in the text format expected by the UI + is_txt_format = (task_data.metadata_output_format.lower() == 'txt') + if is_txt_format: + metadata = {TASK_TEXT_MAPPING[key]: val for key, val in metadata.items() if key in TASK_TEXT_MAPPING} + + entries = [metadata.copy() for _ in range(req.num_outputs)] + for i, entry in enumerate(entries): + entry['Seed' if is_txt_format else 'seed'] = req.seed + i + + return entries + +def get_printable_request(req: GenerateImageRequest): + metadata = req.dict() + del metadata['init_image'] + del metadata['init_image_mask'] + return metadata + +def make_filename_callback(req: GenerateImageRequest, suffix=None, now=None): + if now is None: + now = time.time() + def make_filename(i): + img_id = base64.b64encode(int(now+i).to_bytes(8, 'big')).decode() # Generate unique ID based on time. + img_id = img_id.translate({43:None, 47:None, 61:None})[-8:] # Remove + / = and keep last 8 chars. + + prompt_flattened = filename_regex.sub('_', req.prompt)[:50] + name = f"{prompt_flattened}_{img_id}" + name = name if suffix is None else f'{name}_{suffix}' + return name + + return make_filename diff --git a/ui/index.html b/ui/index.html index 5c76cc88..65f97090 100644 --- a/ui/index.html +++ b/ui/index.html @@ -1,11 +1,12 @@
-