prompt editing

Fixes an issue where it's still possible to write to arbitrary directories through careful use of \.. or /.. in directory patterns

...and fix the regex to work better

reeeegex

.github/ISSUE_TEMPLATE/bug_report.md

@@ -0,0 +1,32 @@
+---
+name: Bug report
+about: Create a report to help us improve
+title: ''
+labels: bug
+assignees: ''
+
+---
+
+**Describe the bug**
+A clear and concise description of what the bug is.
+
+**To Reproduce**
+Steps to reproduce the behavior:
+1. Go to '...'
+2. Click on '....'
+3. Scroll down to '....'
+4. See error
+
+**Expected behavior**
+A clear and concise description of what you expected to happen.
+
+**Screenshots**
+If applicable, add screenshots to help explain your problem.
+
+**Desktop (please complete the following information):**
+ - OS: [e.g. Windows, Linux]
+ - Browser [e.g. chrome, safari]
+ - Commit revision [looks like this: e68484500f76a33ba477d5a99340ab30451e557b; can be seen when launching webui.bat, or obtained manually by running `git rev-parse HEAD`]
+
+**Additional context**
+Add any other context about the problem here.

.github/ISSUE_TEMPLATE/feature_request.md

@@ -0,0 +1,20 @@
+---
+name: Feature request
+about: Suggest an idea for this project
+title: ''
+labels: ''
+assignees: ''
+
+---
+
+**Is your feature request related to a problem? Please describe.**
+A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
+
+**Describe the solution you'd like**
+A clear and concise description of what you want to happen.
+
+**Describe alternatives you've considered**
+A clear and concise description of any alternative solutions or features you've considered.
+
+**Additional context**
+Add any other context or screenshots about the feature request here.

.gitignore

@@ -8,4 +8,11 @@ __pycache__
 /ui-config.json
 /outputs
 /config.json
-/log
+/log
+/webui.settings.bat
+/embeddings
+/styles.csv
+/styles.csv.bak
+/webui-user.bat
+/webui-user.sh
+/interrogate

README.md

@@ -8,7 +8,7 @@ A browser interface based on Gradio library for Stable Diffusion.
 [Detailed feature showcase with images, art by Greg Rutkowski](https://github.com/AUTOMATIC1111/stable-diffusion-webui-feature-showcase)
 
 - Original txt2img and img2img modes
-- One click install and run script (but you still must install python, git and CUDA)
+- One click install and run script (but you still must install python and git)
 - Outpainting
 - Inpainting
 - Prompt matrix
@@ -19,31 +19,34 @@ A browser interface based on Gradio library for Stable Diffusion.
 - Textual Inversion
 - Extras tab with:
   - GFPGAN, neural network that fixes faces
+  - CodeFormer, face restoration tool as an alternative to GFPGAN
   - RealESRGAN, neural network upscaler
   - ESRGAN, neural network with a lot of third party models
 - Resizing aspect ratio options
 - Sampling method selection
 - Interrupt processing at any time
-- 4GB videocard support
+- 4GB video card support
 - Correct seeds for batches
 - Prompt length validation
 - Generation parameters added as text to PNG
 - Tab to view an existing picture's generation parameters
 - Settings page
 - Running custom code from UI
-- Mouseover hints fo most UI elements
+- Mouseover hints for most UI elements
 - Possible to change defaults/mix/max/step values for UI elements via text config
 - Random artist button
 - Tiling support: UI checkbox to create images that can be tiled like textures
 - Progress bar and live image generation preview
+- Negative prompt
+- Styles
+- Variations
+- Seed resizing
+- CLIP interrogator
 
 ## Installing and running
 
 You need [python](https://www.python.org/downloads/windows/) and [git](https://git-scm.com/download/win)
-installed to run this, and an NVidia videocard.
-
-I tested the installation to work Windows with Python 3.8.10, and with Python 3.10.6. You may be able
-to have success with different versions.
+installed to run this, and an NVidia video card.
 
 You need `model.ckpt`, Stable Diffusion model checkpoint, a big file containing the neural network weights. You
 can obtain it from the following places:
@@ -51,65 +54,151 @@ can obtain it from the following places:
  - [file storage](https://drive.yerf.org/wl/?id=EBfTrmcCCUAGaQBXVIj5lJmEhjoP1tgl)
  - magnet:?xt=urn:btih:3a4a612d75ed088ea542acac52f9f45987488d1c&dn=sd-v1-4.ckpt&tr=udp%3a%2f%2ftracker.openbittorrent.com%3a6969%2fannounce&tr=udp%3a%2f%2ftracker.opentrackr.org%3a1337
 
-You optionally can use GPFGAN to improve faces, then you'll need to download the model from [here](https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth).
+You can optionally use GFPGAN to improve faces, to do so you'll need to download the model from [here](https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth) and place it in the same directory as `webui.bat`.
 
 To use ESRGAN models, put them into ESRGAN directory in the same location as webui.py. A file will be loaded
-as model if it has .pth extension. Grab models from the [Model Database](https://upscale.wiki/wiki/Model_Database).
+as a model if it has .pth extension, and it will show up with its name in the UI. Grab models from the [Model Database](https://upscale.wiki/wiki/Model_Database).
+
+> Note: RealESRGAN models are not ESRGAN models, they are not compatible. Do not download RealESRGAN models. Do not place
+RealESRGAN into the directory with ESRGAN models. Thank you.
 
 ### Automatic installation/launch
 
 - install [Python 3.10.6](https://www.python.org/downloads/windows/) and check "Add Python to PATH" during installation. You must install this exact version.
 - install [git](https://git-scm.com/download/win)
-- install [CUDA 11.3](https://developer.nvidia.com/cuda-11.3.0-download-archive?target_os=Windows&target_arch=x86_64)
 - place `model.ckpt` into webui directory, next to `webui.bat`.
 - _*(optional)*_ place `GFPGANv1.3.pth` into webui directory, next to `webui.bat`.
-- run `webui.bat` from Windows Explorer.
+- run `webui-user.bat` from Windows Explorer. Run it as a normal user, ***not*** as administrator.
 
-#### Troublehooting:
+### Running on AMD GPUs
+See the [wiki article](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Running-using-AMD-GPUs) by [cryzed](https://github.com/cryzed).
 
-- According to reports, intallation currently does not work in a directory with spaces in filenames.
-- if your version of Python is not in PATH (or if another version is), edit `webui.bat`, change the line `set PYTHON=python` to say the full path to your python executable: `set PYTHON=B:\soft\Python310\python.exe`. You can do this for python, but not for git.
-- if you get out of memory errors and your videocard has low amount of VRAM (4GB), edit `webui.bat`, change line 5 to from `set COMMANDLINE_ARGS=` to `set COMMANDLINE_ARGS=--medvram` (see below for other possible options)
-- installer creates python virtual environment, so none of installed modules will affect your system installation of python if you had one prior to installing this.
-- to prevent the creation of virtual environment and use your system python, edit `webui.bat` replacing `set VENV_DIR=venv` with `set VENV_DIR=`.
-- webui.bat installs requirements from files `requirements_versions.txt`, which lists versions for modules specifically compatible with Python 3.10.6. If you choose to install for a different version of python, editing `webui.bat` to have `set REQS_FILE=requirements.txt` instead of `set REQS_FILE=requirements_versions.txt` may help (but I still reccomend you to just use the recommended version of python).
+### Linux Automatic installation/launch
+
+Prequisites:
+- For Debian-based:
+```commandline
+sudo apt install wget git python3 python3-venv
+```
+- For Red Hat-based:
+```commandline
+sudo dnf install wget git python3
+```
+
+
+- If you want to install to default directory `/home/$(whoami)/stable-diffusion-webui/`, you can launch directly:
+```commandline
+bash <(wget -qO- https://raw.githubusercontent.com/AUTOMATIC1111/stable-diffusion-webui/master/webui.sh)
+```
+
+
+- If you want to customize the installation just `git clone` the repo where you want it,
+change the variables in `webui-user.sh` and launch in console `bash webui.sh`.
+
+
+
+- place `model.ckpt` into webui directory, next to `webui.py`.
+- _*(optional)*_ place `GFPGANv1.3.pth` into webui directory, next to `webui.py`.
+- run `bash webui.sh`. Run it as a normal user, ***not*** as root.
+
+
+
+#### Troubleshooting
+
+- if your version of Python is not in PATH (or if another version is), edit `webui-user.bat`, and modify the
+line `set PYTHON=python` to say the full path to your python executable, for example: `set PYTHON=B:\soft\Python310\python.exe`.
+You can do this for python, but not for git.
+- if you get out of memory errors and your video-card has a low amount of VRAM (4GB), use custom parameter `set COMMANDLINE_ARGS` (see section below)
+to enable appropriate optimization according to low VRAM guide below (for example, `set COMMANDLINE_ARGS=--medvram --opt-split-attention`).
+- to prevent the creation of virtual environment and use your system python, use custom parameter replacing `set VENV_DIR=-` (see below).
+- webui.bat installs requirements from files `requirements_versions.txt`, which lists versions for modules specifically compatible with
+Python 3.10.6. If you choose to install for a different version of python, using custom parameter `set REQS_FILE=requirements.txt`
+may help (but I still recommend you to just use the recommended version of python).
 - if you feel you broke something and want to reinstall from scratch, delete directories: `venv`, `repositories`.
+- if you get a green or black screen instead of generated pictures, you have a card that doesn't support half precision
+floating point numbers (Known issue with 16xx cards). You must use `--precision full --no-half` in addition to command line
+arguments (set them using `set COMMANDLINE_ARGS`, see below), and the model will take much more space in VRAM (you will likely
+have to also use at least `--medvram`).
+- the installer creates a python virtual environment, so none of the installed modules will affect your system installation of python if
+you had one prior to installing this.
+- About _"You must install this exact version"_ from the instructions above: you can use any version of python you like,
+and it will likely work, but if you want to seek help about things not working, I will not offer help unless you use this
+exact version for my sanity.
 
-### Google collab
+#### How to run with custom parameters
 
-If you don't want or can't run locally, here is google collab that allows you to run the webui:
+It's possible to edit `set COMMANDLINE_ARGS=` line in `webui.bat` to run the program with different command line arguments, but that may lead
+to inconveniences when the file is updated in the repository.
 
-https://colab.research.google.com/drive/1Iy-xW9t1-OQWhb0hNxueGij8phCyluOh
+The recommended way is to use another .bat file named anything you like, set the parameters you want in it, and run webui.bat from it.
+A `webui-user.bat` file included into the repository does exactly this.
 
-### What options to use for low VRAM videocards?
+Here is an example that runs the program with `--opt-split-attention` argument:
+
+```commandline
+@echo off
+
+set COMMANDLINE_ARGS=--opt-split-attention
+
+call webui.bat
+```
+
+Another example, this file will run the program with a custom python path, a different model named `a.ckpt` and without a virtual environment:
+
+```commandline
+@echo off
+
+set PYTHON=b:/soft/Python310/Python.exe
+set VENV_DIR=-
+set COMMANDLINE_ARGS=--ckpt a.ckpt
+
+call webui.bat
+```
+
+### How to create large images?
+Use `--opt-split-attention` parameter. It slows down sampling a tiny bit, but allows you to make gigantic images.
+
+### What options to use for low VRAM video-cards?
+You can, through command line arguments, enable the various optimizations which sacrifice some/a lot of speed in favor of
+using less VRAM. Those arguments are added to the `COMMANDLINE_ARGS` parameter, see section above.
+
+Here's a list of optimization arguments:
 - If you have 4GB VRAM and want to make 512x512 (or maybe up to 640x640) images, use `--medvram`.
 - If you have 4GB VRAM and want to make 512x512 images, but you get an out of memory error with `--medvram`, use `--medvram --opt-split-attention` instead.
 - If you have 4GB VRAM and want to make 512x512 images, and you still get an out of memory error, use `--lowvram --always-batch-cond-uncond --opt-split-attention` instead.
-- If you have 4GB VRAM and want to make images larger than you can with `--medvram`, use  `--lowvram --opt-split-attention`.
-- If you have more VRAM and want to make larger images than you can usually make, use `--medvram --opt-split-attention`. You can use `--lowvram`
+- If you have 4GB VRAM and want to make images larger than you can with `--medvram`, use `--lowvram --opt-split-attention`.
+- If you have more VRAM and want to make larger images than you can usually make (for example 1024x1024 instead of 512x512), use `--medvram --opt-split-attention`. You can use `--lowvram`
 also but the effect will likely be barely noticeable.
 - Otherwise, do not use any of those.
 
-Extra: if you get a green screen instead of generated pictures, you have a card that doesn't support half
-precision floating point numbers. You must use `--precision full --no-half` in addition to other flags,
-and the model will take much more space in VRAM.
-
 ### Running online
 
-Use `--share` option to run online. You will get a xxx.app.gradio link. This is the intended way to use the
-program in collabs.
+Use the `--share` option to run online. You will get a xxx.app.gradio link. This is the intended way to use the
+program in Colab. You may set up authentication for said gradio shared instance with the flag `--gradio-auth username:password`, optionally providing multiple sets of usernames and passwords separated by commas.
 
-Use `--listen` to make the server listen to network connections. This will allow computers on local newtork
+Use `--listen` to make the server listen to network connections. This will allow computers on the local network
 to access the UI, and if you configure port forwarding, also computers on the internet.
 
+Use `--port xxxx` to make the server listen on a specific port, xxxx being the wanted port. Remember that
+all ports below 1024 need root/admin rights, for this reason it is advised to use a port above 1024.
+Defaults to port 7860 if available.
+
+### Google Colab
+
+If you don't want or can't run locally, here is a Google Colab that allows you to run the webui:
+
+https://colab.research.google.com/drive/1Iy-xW9t1-OQWhb0hNxueGij8phCyluOh
+
 ### Textual Inversion
-To make use of pretrained embeddings, create `embeddings` directory (in the same palce as `webui.py`)
-and put your embeddings into it. They must be .pt files, each with only one trained embedding,
-and the filename (without .pt) will be the term you'd use in prompt to get that embedding.
+To make use of pretrained embeddings, create an `embeddings` directory (in the same place as `webui.py`)
+and put your embeddings into it. They must be either .pt or .bin files, each with only one trained embedding,
+and the filename (without .pt/.bin) will be the term you'll use in the prompt to get that embedding.
 
 As an example, I trained one for about 5000 steps: https://files.catbox.moe/e2ui6r.pt; it does not produce
-very good results, but it does work. Download and rename it to Usada Pekora.pt, and put it into embeddings dir
-and use Usada Pekora in prompt.
+very good results, but it does work. To try it out download the file, rename it to `Usada Pekora.pt`, put it into the `embeddings` dir
+and use `Usada Pekora` in the prompt.
+
+You may also try some from the growing library of embeddings at https://huggingface.co/sd-concepts-library, downloading one of the `learned_embeds.bin` files, renaming it to the term you want to use for it in the prompt (be sure to keep the .bin extension) and putting it in your `embeddings` directory.
 
 ### How to change UI defaults?
 
@@ -131,70 +220,90 @@ After running once, a `ui-config.json` file appears in webui directory:
 
 Edit values to your liking and the next time you launch the program they will be applied.
 
+### Almost automatic installation and launch
 
-### Manual instructions
-Alternatively, if you don't want to run webui.bat, here are instructions for installing
-everything by hand:
+Install python and git, place `model.ckpt` and `GFPGANv1.3.pth` into webui directory, run:
 
-```commandline
-:: crate a directory somewhere for stable diffusion and open cmd in it;
-:: make sure you are in the right directory; the command must output the directory you chose
-echo %cd%
+```
+python launch.py
+```
 
-:: install torch with CUDA support. See https://pytorch.org/get-started/locally/ for more instructions if this fails.
+This installs packages via pip. If you need to use a virtual environment, you must set it up yourself. I will not
+provide support for using the web ui this way unless you are using the recommended version of python below.
+
+If you'd like to use command line parameters, use them right there:
+
+```
+python launch.py --opt-split-attention --ckpt ../secret/anime9999.ckpt
+```
+
+### Manual installation
+Alternatively, if you don't want to run the installer, here are instructions for installing
+everything by hand. This can run on both Windows and Linux (if you're on linux, use `ls`
+instead of `dir`). 
+
+```bash
+# install torch with CUDA support. See https://pytorch.org/get-started/locally/ for more instructions if this fails.
 pip install torch --extra-index-url https://download.pytorch.org/whl/cu113
 
-:: check if torch supports GPU; this must output "True". You need CUDA 11. installed for this. You might be able to use
-:: a different version, but this is what I tested.
+# check if torch supports GPU; this must output "True". You need CUDA 11. installed for this. You might be able to use
+# a different version, but this is what I tested.
 python -c "import torch; print(torch.cuda.is_available())"
 
-:: clone Stable Diffusion repositories
-git clone https://github.com/CompVis/stable-diffusion.git
-git clone https://github.com/CompVis/taming-transformers
-
-:: install requirements of Stable Diffusion
-pip install transformers==4.19.2 diffusers invisible-watermark
-
-:: install k-diffusion
-pip install git+https://github.com/crowsonkb/k-diffusion.git
-
-:: (optional) install GFPGAN to fix faces
-pip install git+https://github.com/TencentARC/GFPGAN.git
-
-:: go into stable diffusion's repo directory
-cd stable-diffusion
-
-:: clone web ui
+# clone web ui and go into its directory
 git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git
+cd stable-diffusion-webui
 
-:: install requirements of web ui
-pip install -r stable-diffusion-webui/requirements.txt
+# clone repositories for Stable Diffusion and (optionally) CodeFormer
+mkdir repositories
+git clone https://github.com/CompVis/stable-diffusion.git repositories/stable-diffusion
+git clone https://github.com/CompVis/taming-transformers.git repositories/taming-transformers
+git clone https://github.com/sczhou/CodeFormer.git repositories/CodeFormer
+git clone https://github.com/salesforce/BLIP.git repositories/BLIP
 
-:: update numpy to latest version
-pip install -U numpy
+# install requirements of Stable Diffusion
+pip install transformers==4.19.2 diffusers invisible-watermark --prefer-binary
 
-:: (outside of command line) put stable diffusion model into models/ldm/stable-diffusion-v1/model.ckpt; you'll have
-:: to create one missing directory;
-:: the command below must output something like: 1 File(s) 4,265,380,512 bytes
-dir models\ldm\stable-diffusion-v1\model.ckpt
+# install k-diffusion
+pip install git+https://github.com/crowsonkb/k-diffusion.git --prefer-binary
 
-:: (outside of command line) put the GFPGAN model into same directory as webui script
-:: the command below must output something like: 1 File(s) 348,632,874 bytes
-dir stable-diffusion-webui\GFPGANv1.3.pth
+# (optional) install GFPGAN (face restoration)
+pip install git+https://github.com/TencentARC/GFPGAN.git --prefer-binary
+
+# (optional) install requirements for CodeFormer (face restoration)
+pip install -r repositories/CodeFormer/requirements.txt --prefer-binary
+
+# install requirements of web ui
+pip install -r requirements.txt  --prefer-binary
+
+# update numpy to latest version
+pip install -U numpy  --prefer-binary
+
+# (outside of command line) put stable diffusion model into web ui directory
+# the command below must output something like: 1 File(s) 4,265,380,512 bytes
+dir model.ckpt
+
+# (outside of command line) put the GFPGAN model into web ui directory
+# the command below must output something like: 1 File(s) 348,632,874 bytes
+dir GFPGANv1.3.pth
 ```
 
+> Note: the directory structure for manual instruction has been changed on 2022-09-09 to match automatic installation: previously
+> webui was in a subdirectory of stable diffusion, now it's the reverse. If you followed manual installation before the
+> change, you can still use the program with your existing directory structure.
+
 After that the installation is finished.
 
 Run the command to start web ui:
 
 ```
-python stable-diffusion-webui/webui.py
+python webui.py
 ```
 
 If you have a 4GB video card, run the command with either `--lowvram` or `--medvram` argument:
 
 ```
-python stable-diffusion-webui/webui.py --medvram
+python webui.py --medvram
 ```
 
 After a while, you will get a message like this:
@@ -203,15 +312,57 @@ After a while, you will get a message like this:
 Running on local URL:  http://127.0.0.1:7860/
 ```
 
-Open the URL in browser, and you are good to go.
+Open the URL in a browser, and you are good to go.
 
 
+### Windows 11 WSL2 instructions
+Alternatively, here are instructions for installing under Windows 11 WSL2 Linux distro, everything by hand:
+
+```bash
+# install conda (if not already done)
+wget https://repo.anaconda.com/archive/Anaconda3-2022.05-Linux-x86_64.sh
+chmod +x Anaconda3-2022.05-Linux-x86_64.sh 
+./Anaconda3-2022.05-Linux-x86_64.sh
+
+# Clone webui repo
+git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git
+cd stable-diffusion-webui
+
+# Create and activate conda env
+conda env create -f environment-wsl2.yaml
+conda activate automatic
+
+# (optional) install requirements for GFPGAN (upscaling)
+wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth
+```
+
+After that follow the instructions in the `Manual instructions` section starting at step `:: clone repositories for Stable Diffusion and (optionally) CodeFormer`.
+
+### Custom scripts from users
+
+[A list of custom scripts](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Custom-scripts-from-users), along with installation instructions.
+
+
+
+### img2img alternative test
+- see [this post](https://www.reddit.com/r/StableDiffusion/comments/xboy90/a_better_way_of_doing_img2img_by_finding_the/) on ebaumsworld.com for context.
+- find it in scripts section
+- put description of input image into the Original prompt field
+- use Euler only
+- recommended: 50 steps, low cfg scale between 1 and 2
+- denoising and seed don't matter
+- decode cfg scale between 0 and 1
+- decode steps 50
+- original blue haired woman close nearly reproduces with cfg scale=1.8
+
 ## Credits
 - Stable Diffusion - https://github.com/CompVis/stable-diffusion, https://github.com/CompVis/taming-transformers
 - k-diffusion - https://github.com/crowsonkb/k-diffusion.git
 - GFPGAN - https://github.com/TencentARC/GFPGAN.git
 - ESRGAN - https://github.com/xinntao/ESRGAN
-- Ideas for optimizations and some code (from users) - https://github.com/basujindal/stable-diffusion
+- Ideas for optimizations - https://github.com/basujindal/stable-diffusion
+- Cross Attention layer optimization - https://github.com/Doggettx/stable-diffusion
 - Idea for SD upscale - https://github.com/jquesnelle/txt2imghd
+- CLIP interrogator idea and borrowing some code - https://github.com/pharmapsychotic/clip-interrogator
 - Initial Gradio script - posted on 4chan by an Anonymous user. Thank you Anonymous user.
-- (You)
+- (You)

environment-wsl2.yaml

@@ -0,0 +1,11 @@
+name: automatic
+channels:
+  - pytorch
+  - defaults
+dependencies:
+  - python=3.8.5
+  - pip=20.3
+  - cudatoolkit=11.3
+  - pytorch=1.11.0
+  - torchvision=0.12.0
+  - numpy=1.19.2

launch.py

@@ -0,0 +1,127 @@
+# this scripts installs necessary requirements and launches main program in webui.py
+
+import subprocess
+import os
+import sys
+import importlib.util
+import shlex
+
+dir_repos = "repositories"
+dir_tmp = "tmp"
+
+python = sys.executable
+git = os.environ.get('GIT', "git")
+torch_command = os.environ.get('TORCH_COMMAND', "pip install torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113")
+requirements_file = os.environ.get('REQS_FILE', "requirements_versions.txt")
+commandline_args = os.environ.get('COMMANDLINE_ARGS', "")
+
+k_diffusion_package = os.environ.get('K_DIFFUSION_PACKAGE', "git+https://github.com/crowsonkb/k-diffusion.git@1a0703dfb7d24d8806267c3e7ccc4caf67fd1331")
+gfpgan_package = os.environ.get('GFPGAN_PACKAGE', "git+https://github.com/TencentARC/GFPGAN.git@8d2447a2d918f8eba5a4a01463fd48e45126a379")
+
+stable_diffusion_commit_hash = os.environ.get('STABLE_DIFFUSION_COMMIT_HASH', "69ae4b35e0a0f6ee1af8bb9a5d0016ccb27e36dc")
+taming_transformers_commit_hash = os.environ.get('TAMING_TRANSFORMERS_COMMIT_HASH', "24268930bf1dce879235a7fddd0b2355b84d7ea6")
+codeformer_commit_hash = os.environ.get('CODEFORMER_COMMIT_HASH', "c5b4593074ba6214284d6acd5f1719b6c5d739af")
+blip_commit_hash = os.environ.get('BLIP_COMMIT_HASH', "48211a1594f1321b00f14c9f7a5b4813144b2fb9")
+
+def repo_dir(name):
+    return os.path.join(dir_repos, name)
+
+
+def run(command, desc=None, errdesc=None):
+    if desc is not None:
+        print(desc)
+
+    result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
+
+    if result.returncode != 0:
+
+        message = f"""{errdesc or 'Error running command'}.
+Command: {command}
+Error code: {result.returncode}
+stdout: {result.stdout.decode(encoding="utf8", errors="ignore") if len(result.stdout)>0 else '<empty>'}
+stderr: {result.stderr.decode(encoding="utf8", errors="ignore") if len(result.stderr)>0 else '<empty>'}
+"""
+        raise RuntimeError(message)
+
+    return result.stdout.decode(encoding="utf8", errors="ignore")
+
+
+def run_python(code, desc=None, errdesc=None):
+    return run(f'"{python}" -c "{code}"', desc, errdesc)
+
+
+def run_pip(args, desc=None):
+    return run(f'"{python}" -m pip {args} --prefer-binary', desc=f"Installing {desc}", errdesc=f"Couldn't install {desc}")
+
+
+def check_run(command):
+    result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
+    return result.returncode == 0
+
+
+def check_run_python(code):
+    return check_run(f'"{python}" -c "{code}"')
+
+
+def is_installed(package):
+    try:
+        spec = importlib.util.find_spec(package)
+    except ModuleNotFoundError:
+        return False
+
+    return spec is not None
+
+
+def git_clone(url, dir, name, commithash=None):
+    # TODO clone into temporary dir and move if successful
+
+    if os.path.exists(dir):
+        return
+
+    run(f'"{git}" clone "{url}" "{dir}"', f"Cloning {name} into {dir}...", f"Couldn't clone {name}")
+
+    if commithash is not None:
+        run(f'"{git}" -C {dir} checkout {commithash}', None, "Couldn't checkout {name}'s hash: {commithash}")
+
+
+try:
+    commit = run(f"{git} rev-parse HEAD").strip()
+except Exception:
+    commit = "<none>"
+
+print(f"Python {sys.version}")
+print(f"Commit hash: {commit}")
+
+if not is_installed("torch"):
+    run(f'"{python}" -m {torch_command}', "Installing torch", "Couldn't install torch")
+
+run_python("import torch; assert torch.cuda.is_available(), 'Torch is not able to use GPU'")
+
+if not is_installed("k_diffusion.sampling"):
+    run_pip(f"install {k_diffusion_package}", "k-diffusion")
+
+if not is_installed("gfpgan"):
+    run_pip(f"install {gfpgan_package}", "gfpgan")
+
+os.makedirs(dir_repos, exist_ok=True)
+
+git_clone("https://github.com/CompVis/stable-diffusion.git", repo_dir('stable-diffusion'), "Stable Diffusion", stable_diffusion_commit_hash)
+git_clone("https://github.com/CompVis/taming-transformers.git", repo_dir('taming-transformers'), "Taming Transformers", taming_transformers_commit_hash)
+git_clone("https://github.com/sczhou/CodeFormer.git", repo_dir('CodeFormer'), "CodeFormer", codeformer_commit_hash)
+git_clone("https://github.com/salesforce/BLIP.git", repo_dir('BLIP'), "BLIP", blip_commit_hash)
+
+if not is_installed("lpips"):
+    run_pip(f"install -r {os.path.join(repo_dir('CodeFormer'), 'requirements.txt')}", "requirements for CodeFormer")
+
+run_pip(f"install -r {requirements_file}", "requirements for Web UI")
+
+sys.argv += shlex.split(commandline_args)
+
+
+def start_webui():
+    print(f"Launching Web UI with arguments: {' '.join(sys.argv[1:])}")
+    import webui
+    webui.webui()
+
+
+start_webui()

modules/codeformer_model.py

@@ -1,9 +1,11 @@
 import os
 import sys
 import traceback
+
+import cv2
 import torch
 
-from modules import shared
+from modules import shared, devices
 from modules.paths import script_path
 import modules.shared
 import modules.face_restoration
@@ -49,39 +51,42 @@ def setup_codeformer():
             def create_models(self):
 
                 if self.net is not None and self.face_helper is not None:
+                    self.net.to(devices.device_codeformer)
                     return self.net, self.face_helper
 
-                net = net_class(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, connect_list=['32', '64', '128', '256']).to(shared.device)
+                net = net_class(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, connect_list=['32', '64', '128', '256']).to(devices.device_codeformer)
                 ckpt_path = load_file_from_url(url=pretrain_model_url, model_dir=os.path.join(path, 'weights/CodeFormer'), progress=True)
                 checkpoint = torch.load(ckpt_path)['params_ema']
                 net.load_state_dict(checkpoint)
                 net.eval()
 
-                face_helper = FaceRestoreHelper(1, face_size=512, crop_ratio=(1, 1), det_model='retinaface_resnet50', save_ext='png', use_parse=True, device=shared.device)
+                face_helper = FaceRestoreHelper(1, face_size=512, crop_ratio=(1, 1), det_model='retinaface_resnet50', save_ext='png', use_parse=True, device=devices.device_codeformer)
 
-                if not cmd_opts.unload_gfpgan:
-                    self.net = net
-                    self.face_helper = face_helper
+                self.net = net
+                self.face_helper = face_helper
+                self.net.to(devices.device_codeformer)
 
                 return net, face_helper
 
             def restore(self, np_image, w=None):
                 np_image = np_image[:, :, ::-1]
 
-                net, face_helper = self.create_models()
-                face_helper.clean_all()
-                face_helper.read_image(np_image)
-                face_helper.get_face_landmarks_5(only_center_face=False, resize=640, eye_dist_threshold=5)
-                face_helper.align_warp_face()
+                original_resolution = np_image.shape[0:2]
 
-                for idx, cropped_face in enumerate(face_helper.cropped_faces):
+                self.create_models()
+                self.face_helper.clean_all()
+                self.face_helper.read_image(np_image)
+                self.face_helper.get_face_landmarks_5(only_center_face=False, resize=640, eye_dist_threshold=5)
+                self.face_helper.align_warp_face()
+
+                for idx, cropped_face in enumerate(self.face_helper.cropped_faces):
                     cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
                     normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
-                    cropped_face_t = cropped_face_t.unsqueeze(0).to(shared.device)
+                    cropped_face_t = cropped_face_t.unsqueeze(0).to(devices.device_codeformer)
 
                     try:
                         with torch.no_grad():
-                            output = net(cropped_face_t, w=w if w is not None else shared.opts.code_former_weight, adain=True)[0]
+                            output = self.net(cropped_face_t, w=w if w is not None else shared.opts.code_former_weight, adain=True)[0]
                             restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
                         del output
                         torch.cuda.empty_cache()
@@ -90,12 +95,19 @@ def setup_codeformer():
                         restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))
 
                     restored_face = restored_face.astype('uint8')
-                    face_helper.add_restored_face(restored_face)
+                    self.face_helper.add_restored_face(restored_face)
 
-                face_helper.get_inverse_affine(None)
+                self.face_helper.get_inverse_affine(None)
 
-                restored_img = face_helper.paste_faces_to_input_image()
+                restored_img = self.face_helper.paste_faces_to_input_image()
                 restored_img = restored_img[:, :, ::-1]
+
+                if original_resolution != restored_img.shape[0:2]:
+                    restored_img = cv2.resize(restored_img, (0, 0), fx=original_resolution[1]/restored_img.shape[1], fy=original_resolution[0]/restored_img.shape[0], interpolation=cv2.INTER_LINEAR)
+
+                if shared.opts.face_restoration_unload:
+                    self.net.to(devices.cpu)
+
                 return restored_img
 
         global have_codeformer

modules/devices.py

@@ -0,0 +1,50 @@
+import torch
+
+# has_mps is only available in nightly pytorch (for now), `getattr` for compatibility
+from modules import errors
+
+has_mps = getattr(torch, 'has_mps', False)
+
+cpu = torch.device("cpu")
+
+
+def get_optimal_device():
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+
+    if has_mps:
+        return torch.device("mps")
+
+    return cpu
+
+
+def torch_gc():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+
+
+def enable_tf32():
+    if torch.cuda.is_available():
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+
+
+errors.run(enable_tf32, "Enabling TF32")
+
+
+device = get_optimal_device()
+device_codeformer = cpu if has_mps else device
+
+
+def randn(seed, shape):
+    # Pytorch currently doesn't handle setting randomness correctly when the metal backend is used.
+    if device.type == 'mps':
+        generator = torch.Generator(device=cpu)
+        generator.manual_seed(seed)
+        noise = torch.randn(shape, generator=generator, device=cpu).to(device)
+        return noise
+
+    torch.manual_seed(seed)
+    return torch.randn(shape, device=device)
+

modules/errors.py

@@ -0,0 +1,10 @@
+import sys
+import traceback
+
+
+def run(code, task):
+    try:
+        code()
+    except Exception as e:
+        print(f"{task}: {type(e).__name__}", file=sys.stderr)
+        print(traceback.format_exc(), file=sys.stderr)

modules/esrgan_model.py

@@ -9,19 +9,26 @@ from PIL import Image
 import modules.esrgam_model_arch as arch
 from modules import shared
 from modules.shared import opts
+from modules.devices import has_mps
 import modules.images
 
 
 def load_model(filename):
     # this code is adapted from https://github.com/xinntao/ESRGAN
-
-    pretrained_net = torch.load(filename)
+    pretrained_net = torch.load(filename, map_location='cpu' if has_mps else None)
     crt_model = arch.RRDBNet(3, 3, 64, 23, gc=32)
 
     if 'conv_first.weight' in pretrained_net:
         crt_model.load_state_dict(pretrained_net)
         return crt_model
 
+    if 'model.0.weight' not in pretrained_net:
+        is_realesrgan = "params_ema" in pretrained_net and 'body.0.rdb1.conv1.weight' in pretrained_net["params_ema"]
+        if is_realesrgan:
+            raise Exception("The file is a RealESRGAN model, it can't be used as a ESRGAN model.")
+        else:
+            raise Exception("The file is not a ESRGAN model.")
+
     crt_net = crt_model.state_dict()
     load_net_clean = {}
     for k, v in pretrained_net.items():

modules/extras.py

@@ -0,0 +1,109 @@
+import numpy as np
+from PIL import Image
+
+from modules import processing, shared, images, devices
+from modules.shared import opts
+import modules.gfpgan_model
+from modules.ui import plaintext_to_html
+import modules.codeformer_model
+import piexif
+import piexif.helper
+
+
+cached_images = {}
+
+
+def run_extras(image, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility):
+    devices.torch_gc()
+
+    existing_pnginfo = image.info or {}
+
+    image = image.convert("RGB")
+    info = ""
+
+    outpath = opts.outdir_samples or opts.outdir_extras_samples
+
+    if gfpgan_visibility > 0:
+        restored_img = modules.gfpgan_model.gfpgan_fix_faces(np.array(image, dtype=np.uint8))
+        res = Image.fromarray(restored_img)
+
+        if gfpgan_visibility < 1.0:
+            res = Image.blend(image, res, gfpgan_visibility)
+
+        info += f"GFPGAN visibility:{round(gfpgan_visibility, 2)}\n"
+        image = res
+
+    if codeformer_visibility > 0:
+        restored_img = modules.codeformer_model.codeformer.restore(np.array(image, dtype=np.uint8), w=codeformer_weight)
+        res = Image.fromarray(restored_img)
+
+        if codeformer_visibility < 1.0:
+            res = Image.blend(image, res, codeformer_visibility)
+
+        info += f"CodeFormer w: {round(codeformer_weight, 2)}, CodeFormer visibility:{round(codeformer_visibility)}\n"
+        image = res
+
+    if upscaling_resize != 1.0:
+        def upscale(image, scaler_index, resize):
+            small = image.crop((image.width // 2, image.height // 2, image.width // 2 + 10, image.height // 2 + 10))
+            pixels = tuple(np.array(small).flatten().tolist())
+            key = (resize, scaler_index, image.width, image.height, gfpgan_visibility, codeformer_visibility, codeformer_weight) + pixels
+
+            c = cached_images.get(key)
+            if c is None:
+                upscaler = shared.sd_upscalers[scaler_index]
+                c = upscaler.upscale(image, image.width * resize, image.height * resize)
+                cached_images[key] = c
+
+            return c
+
+        info += f"Upscale: {round(upscaling_resize, 3)}, model:{shared.sd_upscalers[extras_upscaler_1].name}\n"
+        res = upscale(image, extras_upscaler_1, upscaling_resize)
+
+        if extras_upscaler_2 != 0 and extras_upscaler_2_visibility > 0:
+            res2 = upscale(image, extras_upscaler_2, upscaling_resize)
+            info += f"Upscale: {round(upscaling_resize, 3)}, visibility: {round(extras_upscaler_2_visibility, 3)}, model:{shared.sd_upscalers[extras_upscaler_2].name}\n"
+            res = Image.blend(res, res2, extras_upscaler_2_visibility)
+
+        image = res
+
+    while len(cached_images) > 2:
+        del cached_images[next(iter(cached_images.keys()))]
+
+    images.save_image(image, path=outpath, basename="", seed=None, prompt=None, extension=opts.samples_format, info=info, short_filename=True, no_prompt=True, grid=False, pnginfo_section_name="extras", existing_info=existing_pnginfo)
+
+    return image, plaintext_to_html(info), ''
+
+
+def run_pnginfo(image):
+    items = image.info
+
+    if "exif" in image.info:
+        exif = piexif.load(image.info["exif"])
+        exif_comment = (exif or {}).get("Exif", {}).get(piexif.ExifIFD.UserComment, b'')
+        try:
+            exif_comment = piexif.helper.UserComment.load(exif_comment)
+        except ValueError:
+            exif_comment = exif_comment.decode('utf8', errors="ignore")
+
+
+        items['exif comment'] = exif_comment
+
+        for field in ['jfif', 'jfif_version', 'jfif_unit', 'jfif_density', 'dpi', 'exif']:
+            del items[field]
+
+
+    info = ''
+    for key, text in items.items():
+        info += f"""
+<div>
+<p><b>{plaintext_to_html(str(key))}</b></p>
+<p>{plaintext_to_html(str(text))}</p>
+</div>
+""".strip()+"\n"
+
+    if len(info) == 0:
+        message = "Nothing found in the image."
+        info = f"<div><p>{message}<p></div>"
+
+    return '', '', info

modules/gfpgan_model.py

@@ -2,7 +2,7 @@ import os
 import sys
 import traceback
 
-from modules import shared
+from modules import shared, devices
 from modules.shared import cmd_opts
 from modules.paths import script_path
 import modules.face_restoration
@@ -28,24 +28,29 @@ def gfpgan():
     global loaded_gfpgan_model
 
     if loaded_gfpgan_model is not None:
+        loaded_gfpgan_model.gfpgan.to(shared.device)
         return loaded_gfpgan_model
 
     if gfpgan_constructor is None:
         return None
 
     model = gfpgan_constructor(model_path=gfpgan_model_path(), upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None)
-
-    if not cmd_opts.unload_gfpgan:
-        loaded_gfpgan_model = model
+    model.gfpgan.to(shared.device)
+    loaded_gfpgan_model = model
 
     return model
 
 
 def gfpgan_fix_faces(np_image):
+    model = gfpgan()
+
     np_image_bgr = np_image[:, :, ::-1]
-    cropped_faces, restored_faces, gfpgan_output_bgr = gfpgan().enhance(np_image_bgr, has_aligned=False, only_center_face=False, paste_back=True)
+    cropped_faces, restored_faces, gfpgan_output_bgr = model.enhance(np_image_bgr, has_aligned=False, only_center_face=False, paste_back=True)
     np_image = gfpgan_output_bgr[:, :, ::-1]
 
+    if shared.opts.face_restoration_unload:
+        model.gfpgan.to(devices.cpu)
+
     return np_image
 
 

modules/images.py

@@ -1,13 +1,19 @@
+import datetime
 import math
 import os
 from collections import namedtuple
 import re
 
 import numpy as np
+import piexif
+import piexif.helper
 from PIL import Image, ImageFont, ImageDraw, PngImagePlugin
+from fonts.ttf import Roboto
+import string
 
 import modules.shared
-from modules.shared import opts
+from modules import sd_samplers, shared
+from modules.shared import opts, cmd_opts
 
 LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)
 
@@ -132,11 +138,16 @@ def draw_grid_annotations(im, width, height, hor_texts, ver_texts):
 
     fontsize = (width + height) // 25
     line_spacing = fontsize // 2
-    fnt = ImageFont.truetype(opts.font, fontsize)
+
+    try:
+        fnt = ImageFont.truetype(opts.font or Roboto, fontsize)
+    except Exception:
+        fnt = ImageFont.truetype(Roboto, fontsize)
+
     color_active = (0, 0, 0)
     color_inactive = (153, 153, 153)
 
-    pad_left = width * 3 // 4 if len(ver_texts) > 0 else 0
+    pad_left = 0 if sum([sum([len(line.text) for line in lines]) for lines in ver_texts]) == 0 else width * 3 // 4
 
     cols = im.width // width
     rows = im.height // height
@@ -234,47 +245,116 @@ def resize_image(resize_mode, im, width, height):
 
 
 invalid_filename_chars = '<>:"/\\|?*\n'
+re_nonletters = re.compile(r'[\s'+string.punctuation+']+')
 
 
-def sanitize_filename_part(text):
-    return text.replace(' ', '_').translate({ord(x): '' for x in invalid_filename_chars})[:128]
+def sanitize_filename_part(text, replace_spaces=True):
+    if replace_spaces:
+        text = text.replace(' ', '_')
+
+    return text.translate({ord(x): '' for x in invalid_filename_chars})[:128]
 
 
-def save_image(image, path, basename, seed=None, prompt=None, extension='png', info=None, short_filename=False, no_prompt=False):
+def apply_filename_pattern(x, p, seed, prompt):
+    if seed is not None:
+        x = x.replace("[seed]", str(seed))
+    if prompt is not None:
+        x = x.replace("[prompt]", sanitize_filename_part(prompt)[:128])
+        x = x.replace("[prompt_spaces]", sanitize_filename_part(prompt, replace_spaces=False)[:128])
+        if "[prompt_words]" in x:
+            words = [x for x in re_nonletters.split(prompt or "") if len(x) > 0]
+            if len(words) == 0:
+                words = ["empty"]
+
+            x = x.replace("[prompt_words]", " ".join(words[0:8]).strip())
+    if p is not None:
+        x = x.replace("[steps]", str(p.steps))
+        x = x.replace("[cfg]", str(p.cfg_scale))
+        x = x.replace("[width]", str(p.width))
+        x = x.replace("[height]", str(p.height))
+        x = x.replace("[sampler]", sd_samplers.samplers[p.sampler_index].name)
+
+    x = x.replace("[model_hash]", shared.sd_model_hash)
+    x = x.replace("[date]", datetime.date.today().isoformat())
+
+    if cmd_opts.hide_ui_dir_config:
+        x = re.sub(r'^[\\/]+|\.{2,}[\\/]+|[\\/]+\.{2,}', '', x)
+
+    return x
+
+def get_next_sequence_number(path, basename):
+    """
+    Determines and returns the next sequence number to use when saving an image in the specified directory.
+
+    The sequence starts at 0.
+    """
+    result = -1
+    if basename != '':
+        basename = basename + "-"
+
+    prefix_length = len(basename)
+    for p in os.listdir(path):
+        if p.startswith(basename):
+            l = os.path.splitext(p[prefix_length:])[0].split('-') #splits the filename (removing the basename first if one is defined, so the sequence number is always the first element)
+            try:
+                result = max(int(l[0]), result)
+            except ValueError:
+                pass
+
+    return result + 1
+
+def save_image(image, path, basename, seed=None, prompt=None, extension='png', info=None, short_filename=False, no_prompt=False, grid=False, pnginfo_section_name='parameters', p=None, existing_info=None):
     if short_filename or prompt is None or seed is None:
         file_decoration = ""
     elif opts.save_to_dirs:
-        file_decoration = f"-{seed}"
+        file_decoration = opts.samples_filename_pattern or "[seed]"
     else:
-        file_decoration = f"-{seed}-{sanitize_filename_part(prompt)[:128]}"
+        file_decoration = opts.samples_filename_pattern or "[seed]-[prompt_spaces]"
+
+    if file_decoration != "":
+        file_decoration = "-" + file_decoration.lower()
+
+    file_decoration = apply_filename_pattern(file_decoration, p, seed, prompt)
 
     if extension == 'png' and opts.enable_pnginfo and info is not None:
         pnginfo = PngImagePlugin.PngInfo()
-        pnginfo.add_text("parameters", info)
+
+        if existing_info is not None:
+            for k, v in existing_info.items():
+                pnginfo.add_text(k, str(v))
+
+        pnginfo.add_text(pnginfo_section_name, info)
     else:
         pnginfo = None
 
-    if opts.save_to_dirs and not no_prompt:
-        words = re.findall(r'\w+', prompt or "")
-        if len(words) == 0:
-            words = ["empty"]
+    save_to_dirs = (grid and opts.grid_save_to_dirs) or (not grid and opts.save_to_dirs and not no_prompt)
 
-        dirname = " ".join(words[0:opts.save_to_dirs_prompt_len])
+    if save_to_dirs:
+        dirname = apply_filename_pattern(opts.directories_filename_pattern or "[prompt_words]", p, seed, prompt)
         path = os.path.join(path, dirname)
 
     os.makedirs(path, exist_ok=True)
 
-    filecount = len([x for x in os.listdir(path) if os.path.splitext(x)[1] == '.' + extension])
+    basecount = get_next_sequence_number(path, basename)
     fullfn = "a.png"
     fullfn_without_extension = "a"
     for i in range(500):
-        fn = f"{filecount+i:05}" if basename == '' else f"{basename}-{filecount+i:04}"
+        fn = f"{basecount+i:05}" if basename == '' else f"{basename}-{basecount+i:04}"
         fullfn = os.path.join(path, f"{fn}{file_decoration}.{extension}")
         fullfn_without_extension = os.path.join(path, f"{fn}{file_decoration}")
         if not os.path.exists(fullfn):
             break
 
-    image.save(fullfn, quality=opts.jpeg_quality, pnginfo=pnginfo)
+    if extension.lower() in ("jpg", "jpeg", "webp"):
+        exif_bytes = piexif.dump({
+            "Exif": {
+                piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(info, encoding="unicode")
+            },
+        })
+    else:
+        exif_bytes = None
+
+    image.save(fullfn, quality=opts.jpeg_quality, pnginfo=pnginfo, exif=exif_bytes)
 
     target_side_length = 4000
     oversize = image.width > target_side_length or image.height > target_side_length
@@ -286,7 +366,7 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
         elif oversize:
             image = image.resize((image.width * target_side_length // image.height, target_side_length), LANCZOS)
 
-        image.save(f"{fullfn_without_extension}.jpg", quality=opts.jpeg_quality, pnginfo=pnginfo)
+        image.save(fullfn, quality=opts.jpeg_quality, exif=exif_bytes)
 
     if opts.save_txt and info is not None:
         with open(f"{fullfn_without_extension}.txt", "w", encoding="utf8") as file:
@@ -307,7 +387,7 @@ class Upscaler:
             img = self.do_upscale(img)
 
         if img.width != w or img.height != h:
-            img = img.resize((w, h), resample=LANCZOS)
+            img = img.resize((int(w), int(h)), resample=LANCZOS)
 
         return img
 

modules/img2img.py

@@ -1,6 +1,8 @@
 import math
-from PIL import Image
+import numpy as np
+from PIL import Image, ImageOps, ImageChops
 
+from modules import devices
 from modules.processing import Processed, StableDiffusionProcessingImg2Img, process_images
 from modules.shared import opts, state
 import modules.shared as shared
@@ -9,14 +11,20 @@ from modules.ui import plaintext_to_html
 import modules.images as images
 import modules.scripts
 
-def img2img(prompt: str, init_img, init_img_with_mask, steps: int, sampler_index: int, mask_blur: int, inpainting_fill: int, restore_faces: bool, tiling: bool, mode: int, n_iter: int, batch_size: int, cfg_scale: float, denoising_strength: float, seed: int, height: int, width: int, resize_mode: int, upscaler_index: str, upscale_overlap: int, inpaint_full_res: bool, inpainting_mask_invert: int, *args):
+def img2img(prompt: str, negative_prompt: str, prompt_style: str, prompt_style2: str, init_img, init_img_with_mask, init_mask, mask_mode, steps: int, sampler_index: int, mask_blur: int, inpainting_fill: int, restore_faces: bool, tiling: bool, mode: int, n_iter: int, batch_size: int, cfg_scale: float, denoising_strength: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, height: int, width: int, resize_mode: int, upscaler_index: str, upscale_overlap: int, inpaint_full_res: bool, inpainting_mask_invert: int, *args):
     is_inpaint = mode == 1
-    is_loopback = mode == 2
-    is_upscale = mode == 3
+    is_upscale = mode == 2
 
     if is_inpaint:
-        image = init_img_with_mask['image']
-        mask = init_img_with_mask['mask']
+        if mask_mode == 0:
+            image = init_img_with_mask['image']
+            mask = init_img_with_mask['mask']
+            alpha_mask = ImageOps.invert(image.split()[-1]).convert('L').point(lambda x: 255 if x > 0 else 0, mode='1')
+            mask = ImageChops.lighter(alpha_mask, mask.convert('L')).convert('L')
+            image = image.convert('RGB')
+        else:
+            image = init_img
+            mask = init_mask
     else:
         image = init_img
         mask = None
@@ -28,7 +36,13 @@ def img2img(prompt: str, init_img, init_img_with_mask, steps: int, sampler_index
         outpath_samples=opts.outdir_samples or opts.outdir_img2img_samples,
         outpath_grids=opts.outdir_grids or opts.outdir_img2img_grids,
         prompt=prompt,
+        negative_prompt=negative_prompt,
+        styles=[prompt_style, prompt_style2],
         seed=seed,
+        subseed=subseed,
+        subseed_strength=subseed_strength,
+        seed_resize_from_h=seed_resize_from_h,
+        seed_resize_from_w=seed_resize_from_w,
         sampler_index=sampler_index,
         batch_size=batch_size,
         n_iter=n_iter,
@@ -46,93 +60,69 @@ def img2img(prompt: str, init_img, init_img_with_mask, steps: int, sampler_index
         denoising_strength=denoising_strength,
         inpaint_full_res=inpaint_full_res,
         inpainting_mask_invert=inpainting_mask_invert,
-        extra_generation_params={"Denoising Strength": denoising_strength}
     )
+    print(f"\nimg2img: {prompt}", file=shared.progress_print_out)
 
-    if is_loopback:
-        output_images, info = None, None
-        history = []
-        initial_seed = None
+    if is_upscale:
         initial_info = None
 
-        state.job_count = n_iter
-
-        for i in range(n_iter):
-
-            p.n_iter = 1
-            p.batch_size = 1
-            p.do_not_save_grid = True
-
-            state.job = f"Batch {i + 1} out of {n_iter}"
-            processed = process_images(p)
-
-            if initial_seed is None:
-                initial_seed = processed.seed
-                initial_info = processed.info
-
-            p.init_images = [processed.images[0]]
-            p.seed = processed.seed + 1
-            p.denoising_strength = max(p.denoising_strength * 0.95, 0.1)
-            history.append(processed.images[0])
-
-        grid = images.image_grid(history, batch_size, rows=1)
-
-        images.save_image(grid, p.outpath_grids, "grid", initial_seed, prompt, opts.grid_format, info=info, short_filename=not opts.grid_extended_filename)
-
-        processed = Processed(p, history, initial_seed, initial_info)
-
-    elif is_upscale:
-        initial_seed = None
-        initial_info = None
+        processing.fix_seed(p)
+        seed = p.seed
 
         upscaler = shared.sd_upscalers[upscaler_index]
         img = upscaler.upscale(init_img, init_img.width * 2, init_img.height * 2)
 
-        processing.torch_gc()
+        devices.torch_gc()
 
         grid = images.split_grid(img, tile_w=width, tile_h=height, overlap=upscale_overlap)
 
+        upscale_count = p.n_iter
         p.n_iter = 1
         p.do_not_save_grid = True
         p.do_not_save_samples = True
 
         work = []
-        work_results = []
 
         for y, h, row in grid.tiles:
             for tiledata in row:
                 work.append(tiledata[2])
 
         batch_count = math.ceil(len(work) / p.batch_size)
-        print(f"SD upscaling will process a total of {len(work)} images tiled as {len(grid.tiles[0][2])}x{len(grid.tiles)} in a total of {batch_count} batches.")
+        state.job_count = batch_count * upscale_count
 
-        state.job_count = batch_count
+        print(f"SD upscaling will process a total of {len(work)} images tiled as {len(grid.tiles[0][2])}x{len(grid.tiles)} per upscale in a total of {state.job_count} batches.")
 
-        for i in range(batch_count):
-            p.init_images = work[i*p.batch_size:(i+1)*p.batch_size]
+        result_images = []
+        for n in range(upscale_count):
+            start_seed = seed + n
+            p.seed = start_seed
 
-            state.job = f"Batch {i + 1} out of {batch_count}"
-            processed = process_images(p)
+            work_results = []
+            for i in range(batch_count):
+                p.init_images = work[i*p.batch_size:(i+1)*p.batch_size]
 
-            if initial_seed is None:
-                initial_seed = processed.seed
-                initial_info = processed.info
+                state.job = f"Batch {i + 1} out of {state.job_count}"
+                processed = process_images(p)
 
-            p.seed = processed.seed + 1
-            work_results += processed.images
+                if initial_info is None:
+                    initial_info = processed.info
 
-        image_index = 0
-        for y, h, row in grid.tiles:
-            for tiledata in row:
-                tiledata[2] = work_results[image_index] if image_index < len(work_results) else Image.new("RGB", (p.width, p.height))
-                image_index += 1
+                p.seed = processed.seed + 1
+                work_results += processed.images
 
-        combined_image = images.combine_grid(grid)
+            image_index = 0
+            for y, h, row in grid.tiles:
+                for tiledata in row:
+                    tiledata[2] = work_results[image_index] if image_index < len(work_results) else Image.new("RGB", (p.width, p.height))
+                    image_index += 1
 
-        if opts.samples_save:
-            images.save_image(combined_image, p.outpath_samples, "", initial_seed, prompt, opts.grid_format, info=initial_info)
+            combined_image = images.combine_grid(grid)
+            result_images.append(combined_image)
 
-        processed = Processed(p, [combined_image], initial_seed, initial_info)
+            if opts.samples_save:
+                images.save_image(combined_image, p.outpath_samples, "", start_seed, prompt, opts.samples_format, info=initial_info, p=p)
+
+        processed = Processed(p, result_images, seed, initial_info)
 
     else:
 
@@ -141,5 +131,6 @@ def img2img(prompt: str, init_img, init_img_with_mask, steps: int, sampler_index
         if processed is None:
             processed = process_images(p)
 
+    shared.total_tqdm.clear()
 
     return processed.images, processed.js(), plaintext_to_html(processed.info)

modules/interrogate.py

@@ -0,0 +1,167 @@
+import contextlib
+import os
+import sys
+import traceback
+from collections import namedtuple
+import re
+
+import torch
+
+from torchvision import transforms
+from torchvision.transforms.functional import InterpolationMode
+
+import modules.shared as shared
+from modules import devices, paths, lowvram
+
+blip_image_eval_size = 384
+blip_model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_caption_capfilt_large.pth'
+clip_model_name = 'ViT-L/14'
+
+Category = namedtuple("Category", ["name", "topn", "items"])
+
+re_topn = re.compile(r"\.top(\d+)\.")
+
+class InterrogateModels:
+    blip_model = None
+    clip_model = None
+    clip_preprocess = None
+    categories = None
+    dtype = None
+
+    def __init__(self, content_dir):
+        self.categories = []
+
+        if os.path.exists(content_dir):
+            for filename in os.listdir(content_dir):
+                m = re_topn.search(filename)
+                topn = 1 if m is None else int(m.group(1))
+
+                with open(os.path.join(content_dir, filename), "r", encoding="utf8") as file:
+                    lines = [x.strip() for x in file.readlines()]
+
+                self.categories.append(Category(name=filename, topn=topn, items=lines))
+
+    def load_blip_model(self):
+        import models.blip
+
+        blip_model = models.blip.blip_decoder(pretrained=blip_model_url, image_size=blip_image_eval_size, vit='base', med_config=os.path.join(paths.paths["BLIP"], "configs", "med_config.json"))
+        blip_model.eval()
+
+        return blip_model
+
+    def load_clip_model(self):
+        import clip
+
+        model, preprocess = clip.load(clip_model_name)
+        model.eval()
+        model = model.to(shared.device)
+
+        return model, preprocess
+
+    def load(self):
+        if self.blip_model is None:
+            self.blip_model = self.load_blip_model()
+            if not shared.cmd_opts.no_half:
+                self.blip_model = self.blip_model.half()
+
+        self.blip_model = self.blip_model.to(shared.device)
+
+        if self.clip_model is None:
+            self.clip_model, self.clip_preprocess = self.load_clip_model()
+            if not shared.cmd_opts.no_half:
+                self.clip_model = self.clip_model.half()
+
+        self.clip_model = self.clip_model.to(shared.device)
+
+        self.dtype = next(self.clip_model.parameters()).dtype
+
+    def send_clip_to_ram(self):
+        if not shared.opts.interrogate_keep_models_in_memory:
+            if self.clip_model is not None:
+                self.clip_model = self.clip_model.to(devices.cpu)
+
+    def send_blip_to_ram(self):
+        if not shared.opts.interrogate_keep_models_in_memory:
+            if self.blip_model is not None:
+                self.blip_model = self.blip_model.to(devices.cpu)
+
+    def unload(self):
+        self.send_clip_to_ram()
+        self.send_blip_to_ram()
+
+        devices.torch_gc()
+
+    def rank(self, image_features, text_array, top_count=1):
+        import clip
+
+        if shared.opts.interrogate_clip_dict_limit != 0:
+            text_array = text_array[0:int(shared.opts.interrogate_clip_dict_limit)]
+
+        top_count = min(top_count, len(text_array))
+        text_tokens = clip.tokenize([text for text in text_array]).to(shared.device)
+        text_features = self.clip_model.encode_text(text_tokens).type(self.dtype)
+        text_features /= text_features.norm(dim=-1, keepdim=True)
+
+        similarity = torch.zeros((1, len(text_array))).to(shared.device)
+        for i in range(image_features.shape[0]):
+            similarity += (100.0 * image_features[i].unsqueeze(0) @ text_features.T).softmax(dim=-1)
+        similarity /= image_features.shape[0]
+
+        top_probs, top_labels = similarity.cpu().topk(top_count, dim=-1)
+        return [(text_array[top_labels[0][i].numpy()], (top_probs[0][i].numpy()*100)) for i in range(top_count)]
+
+    def generate_caption(self, pil_image):
+        gpu_image = transforms.Compose([
+            transforms.Resize((blip_image_eval_size, blip_image_eval_size), interpolation=InterpolationMode.BICUBIC),
+            transforms.ToTensor(),
+            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
+        ])(pil_image).unsqueeze(0).type(self.dtype).to(shared.device)
+
+        with torch.no_grad():
+            caption = self.blip_model.generate(gpu_image, sample=False, num_beams=shared.opts.interrogate_clip_num_beams, min_length=shared.opts.interrogate_clip_min_length, max_length=shared.opts.interrogate_clip_max_length)
+
+        return caption[0]
+
+    def interrogate(self, pil_image):
+        res = None
+
+        try:
+
+            if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
+                lowvram.send_everything_to_cpu()
+                devices.torch_gc()
+
+            self.load()
+
+            caption = self.generate_caption(pil_image)
+            self.send_blip_to_ram()
+            devices.torch_gc()
+
+            res = caption
+
+            cilp_image = self.clip_preprocess(pil_image).unsqueeze(0).type(self.dtype).to(shared.device)
+
+            precision_scope = torch.autocast if shared.cmd_opts.precision == "autocast" else contextlib.nullcontext
+            with torch.no_grad(), precision_scope("cuda"):
+                image_features = self.clip_model.encode_image(cilp_image).type(self.dtype)
+
+                image_features /= image_features.norm(dim=-1, keepdim=True)
+
+                if shared.opts.interrogate_use_builtin_artists:
+                    artist = self.rank(image_features, ["by " + artist.name for artist in shared.artist_db.artists])[0]
+
+                    res += ", " + artist[0]
+
+                for name, topn, items in self.categories:
+                    matches = self.rank(image_features, items, top_count=topn)
+                    for match, score in matches:
+                        res += ", " + match
+
+        except Exception:
+            print(f"Error interrogating", file=sys.stderr)
+            print(traceback.format_exc(), file=sys.stderr)
+            res += "<error>"
+
+        self.unload()
+
+        return res

modules/lowvram.py

@@ -1,9 +1,18 @@
 import torch
+from modules.devices import get_optimal_device
 
 module_in_gpu = None
 cpu = torch.device("cpu")
-gpu = torch.device("cuda")
-device = gpu if torch.cuda.is_available() else cpu
+device = gpu = get_optimal_device()
+
+
+def send_everything_to_cpu():
+    global module_in_gpu
+
+    if module_in_gpu is not None:
+        module_in_gpu.to(cpu)
+
+    module_in_gpu = None
 
 
 def setup_for_low_vram(sd_model, use_medvram):

modules/paths.py

@@ -7,17 +7,18 @@ sys.path.insert(0, script_path)
 
 # search for directory of stable diffsuion in following palces
 sd_path = None
-possible_sd_paths = ['.', os.path.dirname(script_path), os.path.join(script_path, 'repositories/stable-diffusion')]
+possible_sd_paths = [os.path.join(script_path, 'repositories/stable-diffusion'), '.', os.path.dirname(script_path)]
 for possible_sd_path in possible_sd_paths:
     if os.path.exists(os.path.join(possible_sd_path, 'ldm/models/diffusion/ddpm.py')):
         sd_path = os.path.abspath(possible_sd_path)
 
-assert sd_path is not None, "Couldn't find Stable Diffusion in any of: " + possible_sd_paths
+assert sd_path is not None, "Couldn't find Stable Diffusion in any of: " + str(possible_sd_paths)
 
 path_dirs = [
     (sd_path, 'ldm', 'Stable Diffusion'),
     (os.path.join(sd_path, '../taming-transformers'), 'taming', 'Taming Transformers'),
     (os.path.join(sd_path, '../CodeFormer'), 'inference_codeformer.py', 'CodeFormer'),
+    (os.path.join(sd_path, '../BLIP'), 'models/blip.py', 'BLIP'),
 ]
 
 paths = {}

modules/processing.py

@@ -8,35 +8,57 @@ import torch
 import numpy as np
 from PIL import Image, ImageFilter, ImageOps
 import random
+import cv2
+from skimage import exposure
 
 import modules.sd_hijack
+from modules import devices, prompt_parser
 from modules.sd_hijack import model_hijack
 from modules.sd_samplers import samplers, samplers_for_img2img
 from modules.shared import opts, cmd_opts, state
 import modules.shared as shared
 import modules.face_restoration
 import modules.images as images
+import modules.styles
+
 
 # some of those options should not be changed at all because they would break the model, so I removed them from options.
 opt_C = 4
 opt_f = 8
 
 
-def torch_gc():
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.ipc_collect()
+def setup_color_correction(image):
+    correction_target = cv2.cvtColor(np.asarray(image.copy()), cv2.COLOR_RGB2LAB)
+    return correction_target
+
+
+def apply_color_correction(correction, image):
+    image = Image.fromarray(cv2.cvtColor(exposure.match_histograms(
+        cv2.cvtColor(
+            np.asarray(image),
+            cv2.COLOR_RGB2LAB
+        ),
+        correction,
+        channel_axis=2
+    ), cv2.COLOR_LAB2RGB).astype("uint8"))
+
+    return image
 
 
 class StableDiffusionProcessing:
-    def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt="", seed=-1, sampler_index=0, batch_size=1, n_iter=1, steps=50, cfg_scale=7.0, width=512, height=512, restore_faces=False, tiling=False, do_not_save_samples=False, do_not_save_grid=False, extra_generation_params=None, overlay_images=None, negative_prompt=None):
+    def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt="", styles=None, seed=-1, subseed=-1, subseed_strength=0, seed_resize_from_h=-1, seed_resize_from_w=-1, sampler_index=0, batch_size=1, n_iter=1, steps=50, cfg_scale=7.0, width=512, height=512, restore_faces=False, tiling=False, do_not_save_samples=False, do_not_save_grid=False, extra_generation_params=None, overlay_images=None, negative_prompt=None):
         self.sd_model = sd_model
         self.outpath_samples: str = outpath_samples
         self.outpath_grids: str = outpath_grids
         self.prompt: str = prompt
         self.prompt_for_display: str = None
         self.negative_prompt: str = (negative_prompt or "")
+        self.styles: str = styles
         self.seed: int = seed
+        self.subseed: int = subseed
+        self.subseed_strength: float = subseed_strength
+        self.seed_resize_from_h: int = seed_resize_from_h
+        self.seed_resize_from_w: int = seed_resize_from_w
         self.sampler_index: int = sampler_index
         self.batch_size: int = batch_size
         self.n_iter: int = n_iter
@@ -51,8 +73,9 @@ class StableDiffusionProcessing:
         self.extra_generation_params: dict = extra_generation_params
         self.overlay_images = overlay_images
         self.paste_to = None
+        self.color_corrections = None
 
-    def init(self):
+    def init(self, seed):
         pass
 
     def sample(self, x, conditioning, unconditional_conditioning):
@@ -63,6 +86,7 @@ class Processed:
     def __init__(self, p: StableDiffusionProcessing, images_list, seed, info):
         self.images = images_list
         self.prompt = p.prompt
+        self.negative_prompt = p.negative_prompt
         self.seed = seed
         self.info = info
         self.width = p.width
@@ -74,6 +98,7 @@ class Processed:
     def js(self):
         obj = {
             "prompt": self.prompt if type(self.prompt) != list else self.prompt[0],
+            "negative_prompt": self.negative_prompt if type(self.negative_prompt) != list else self.negative_prompt[0],
             "seed": int(self.seed if type(self.seed) != list else self.seed[0]),
             "width": self.width,
             "height": self.height,
@@ -84,69 +109,123 @@ class Processed:
 
         return json.dumps(obj)
 
+# from https://discuss.pytorch.org/t/help-regarding-slerp-function-for-generative-model-sampling/32475/3
+def slerp(val, low, high):
+    low_norm = low/torch.norm(low, dim=1, keepdim=True)
+    high_norm = high/torch.norm(high, dim=1, keepdim=True)
+    omega = torch.acos((low_norm*high_norm).sum(1))
+    so = torch.sin(omega)
+    res = (torch.sin((1.0-val)*omega)/so).unsqueeze(1)*low + (torch.sin(val*omega)/so).unsqueeze(1) * high
+    return res
 
-def create_random_tensors(shape, seeds):
+
+def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, seed_resize_from_h=0, seed_resize_from_w=0):
     xs = []
-    for seed in seeds:
-        torch.manual_seed(seed)
+    for i, seed in enumerate(seeds):
+        noise_shape = shape if seed_resize_from_h <= 0 or seed_resize_from_w <= 0 else (shape[0], seed_resize_from_h//8, seed_resize_from_w//8)
+
+        subnoise = None
+        if subseeds is not None:
+            subseed = 0 if i >= len(subseeds) else subseeds[i]
+
+            subnoise = devices.randn(subseed, noise_shape)
 
         # randn results depend on device; gpu and cpu get different results for same seed;
         # the way I see it, it's better to do this on CPU, so that everyone gets same result;
-        # but the original script had it like this so I do not dare change it for now because
+        # but the original script had it like this, so I do not dare change it for now because
         # it will break everyone's seeds.
-        xs.append(torch.randn(shape, device=shared.device))
-    x = torch.stack(xs)
+        noise = devices.randn(seed, noise_shape)
+
+        if subnoise is not None:
+            #noise = subnoise * subseed_strength + noise * (1 - subseed_strength)
+            noise = slerp(subseed_strength, noise, subnoise)
+
+        if noise_shape != shape:
+            #noise = torch.nn.functional.interpolate(noise.unsqueeze(1), size=shape[1:], mode="bilinear").squeeze()
+            x = devices.randn(seed, shape)
+            dx = (shape[2] - noise_shape[2]) // 2
+            dy = (shape[1] - noise_shape[1]) // 2
+            w = noise_shape[2] if dx >= 0 else noise_shape[2] + 2 * dx
+            h = noise_shape[1] if dy >= 0 else noise_shape[1] + 2 * dy
+            tx = 0 if dx < 0 else dx
+            ty = 0 if dy < 0 else dy
+            dx = max(-dx, 0)
+            dy = max(-dy, 0)
+
+            x[:, ty:ty+h, tx:tx+w] = noise[:, dy:dy+h, dx:dx+w]
+            noise = x
+
+
+
+        xs.append(noise)
+    x = torch.stack(xs).to(shared.device)
     return x
 
 
-def set_seed(seed):
-    return int(random.randrange(4294967294)) if seed is None or seed == -1 else seed
+def fix_seed(p):
+    p.seed = int(random.randrange(4294967294)) if p.seed is None or p.seed == -1 else p.seed
+    p.subseed = int(random.randrange(4294967294)) if p.subseed is None or p.subseed == -1 else p.subseed
 
 
 def process_images(p: StableDiffusionProcessing) -> Processed:
     """this is the main loop that both txt2img and img2img use; it calls func_init once inside all the scopes and func_sample once per batch"""
 
-    prompt = p.prompt
-
     assert p.prompt is not None
-    torch_gc()
+    devices.torch_gc()
 
-    seed = set_seed(p.seed)
+    fix_seed(p)
 
     os.makedirs(p.outpath_samples, exist_ok=True)
     os.makedirs(p.outpath_grids, exist_ok=True)
 
     modules.sd_hijack.model_hijack.apply_circular(p.tiling)
 
-    comments = []
+    comments = {}
 
-    if type(prompt) == list:
-        all_prompts = prompt
-    else:
-        all_prompts = p.batch_size * p.n_iter * [prompt]
+    shared.prompt_styles.apply_styles(p)
 
-    if type(seed) == list:
-        all_seeds = seed
+    if type(p.prompt) == list:
+        all_prompts = p.prompt
     else:
-        all_seeds = [int(seed + x) for x in range(len(all_prompts))]
+        all_prompts = p.batch_size * p.n_iter * [p.prompt]
+
+    if type(p.seed) == list:
+        all_seeds = p.seed
+    else:
+        all_seeds = [int(p.seed + (x if p.subseed_strength == 0 else 0)) for x in range(len(all_prompts))]
+
+    if type(p.subseed) == list:
+        all_subseeds = p.subseed
+    else:
+        all_subseeds = [int(p.subseed + x) for x in range(len(all_prompts))]
 
     def infotext(iteration=0, position_in_batch=0):
+        index = position_in_batch + iteration * p.batch_size
+
         generation_params = {
             "Steps": p.steps,
             "Sampler": samplers[p.sampler_index].name,
             "CFG scale": p.cfg_scale,
-            "Seed": all_seeds[position_in_batch + iteration * p.batch_size],
+            "Seed": all_seeds[index],
             "Face restoration": (opts.face_restoration_model if p.restore_faces else None),
+            "Size": f"{p.width}x{p.height}",
+            "Model hash": (None if not opts.add_model_hash_to_info or not shared.sd_model_hash else shared.sd_model_hash),
             "Batch size": (None if p.batch_size < 2 else p.batch_size),
             "Batch pos": (None if p.batch_size < 2 else position_in_batch),
+            "Variation seed": (None if p.subseed_strength == 0 else all_subseeds[index]),
+            "Variation seed strength": (None if p.subseed_strength == 0 else p.subseed_strength),
+            "Seed resize from": (None if p.seed_resize_from_w == 0 or p.seed_resize_from_h == 0 else f"{p.seed_resize_from_w}x{p.seed_resize_from_h}"),
+            "Denoising strength": getattr(p, 'denoising_strength', None),
         }
 
         if p.extra_generation_params is not None:
             generation_params.update(p.extra_generation_params)
 
         generation_params_text = ", ".join([k if k == v else f'{k}: {v}' for k, v in generation_params.items() if v is not None])
+        
+        negative_prompt_text = "\nNegative prompt: " + p.negative_prompt if p.negative_prompt else ""
 
-        return f"{p.prompt_for_display or prompt}\n{generation_params_text}".strip() + "".join(["\n\n" + x for x in comments])
+        return f"{all_prompts[index]}{negative_prompt_text}\n{generation_params_text}".strip() + "".join(["\n\n" + x for x in comments])
 
     if os.path.exists(cmd_opts.embeddings_dir):
         model_hijack.load_textual_inversion_embeddings(cmd_opts.embeddings_dir, p.sd_model)
@@ -155,7 +234,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
     precision_scope = torch.autocast if cmd_opts.precision == "autocast" else contextlib.nullcontext
     ema_scope = (contextlib.nullcontext if cmd_opts.lowvram else p.sd_model.ema_scope)
     with torch.no_grad(), precision_scope("cuda"), ema_scope():
-        p.init()
+        p.init(seed=all_seeds[0])
 
         if state.job_count == -1:
             state.job_count = p.n_iter
@@ -166,15 +245,19 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
 
             prompts = all_prompts[n * p.batch_size:(n + 1) * p.batch_size]
             seeds = all_seeds[n * p.batch_size:(n + 1) * p.batch_size]
+            subseeds = all_subseeds[n * p.batch_size:(n + 1) * p.batch_size]
 
-            uc = p.sd_model.get_learned_conditioning(len(prompts) * [p.negative_prompt])
-            c = p.sd_model.get_learned_conditioning(prompts)
+            #uc = p.sd_model.get_learned_conditioning(len(prompts) * [p.negative_prompt])
+            #c = p.sd_model.get_learned_conditioning(prompts)
+            uc = prompt_parser.get_learned_conditioning(len(prompts) * [p.negative_prompt], p.steps)
+            c = prompt_parser.get_learned_conditioning(prompts, p.steps)
 
             if len(model_hijack.comments) > 0:
-                comments += model_hijack.comments
+                for comment in model_hijack.comments:
+                    comments[comment] = 1
 
             # we manually generate all input noises because each one should have a specific seed
-            x = create_random_tensors([opt_C, p.height // opt_f, p.width // opt_f], seeds=seeds)
+            x = create_random_tensors([opt_C, p.height // opt_f, p.width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=p.subseed_strength, seed_resize_from_h=p.seed_resize_from_h, seed_resize_from_w=p.seed_resize_from_w)
 
             if p.n_iter > 1:
                 shared.state.job = f"Batch {n+1} out of {p.n_iter}"
@@ -189,17 +272,27 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
             x_samples_ddim = p.sd_model.decode_first_stage(samples_ddim)
             x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
 
+            if opts.filter_nsfw:
+                import modules.safety as safety
+                x_samples_ddim = modules.safety.censor_batch(x_samples_ddim)
+
             for i, x_sample in enumerate(x_samples_ddim):
                 x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)
                 x_sample = x_sample.astype(np.uint8)
 
                 if p.restore_faces:
-                    torch_gc()
+                    if opts.save and not p.do_not_save_samples and opts.save_images_before_face_restoration:
+                        images.save_image(Image.fromarray(x_sample), p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i), p=p)
+
+                    devices.torch_gc()
 
                     x_sample = modules.face_restoration.restore_faces(x_sample)
 
                 image = Image.fromarray(x_sample)
 
+                if p.color_corrections is not None and i < len(p.color_corrections):
+                    image = apply_color_correction(p.color_corrections[i], image)
+
                 if p.overlay_images is not None and i < len(p.overlay_images):
                     overlay = p.overlay_images[i]
 
@@ -215,32 +308,30 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
                     image = image.convert('RGB')
 
                 if opts.samples_save and not p.do_not_save_samples:
-                    images.save_image(image, p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i))
+                    images.save_image(image, p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i), p=p)
 
                 output_images.append(image)
 
             state.nextjob()
 
         unwanted_grid_because_of_img_count = len(output_images) < 2 and opts.grid_only_if_multiple
-        if not p.do_not_save_grid and not unwanted_grid_because_of_img_count:
-            return_grid = opts.return_grid
-
+        if (opts.return_grid or opts.grid_save) and not p.do_not_save_grid and not unwanted_grid_because_of_img_count:
             grid = images.image_grid(output_images, p.batch_size)
 
-            if return_grid:
+            if opts.return_grid:
                 output_images.insert(0, grid)
 
             if opts.grid_save:
-                images.save_image(grid, p.outpath_grids, "grid", seed, all_prompts[0], opts.grid_format, info=infotext(), short_filename=not opts.grid_extended_filename)
+                images.save_image(grid, p.outpath_grids, "grid", all_seeds[0], all_prompts[0], opts.grid_format, info=infotext(), short_filename=not opts.grid_extended_filename, p=p)
 
-    torch_gc()
-    return Processed(p, output_images, seed, infotext())
+    devices.torch_gc()
+    return Processed(p, output_images, all_seeds[0], infotext())
 
 
 class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
     sampler = None
 
-    def init(self):
+    def init(self, seed):
         self.sampler = samplers[self.sampler_index].constructor(self.sd_model)
 
     def sample(self, x, conditioning, unconditional_conditioning):
@@ -320,7 +411,7 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
         self.mask = None
         self.nmask = None
 
-    def init(self):
+    def init(self, seed):
         self.sampler = samplers_for_img2img[self.sampler_index].constructor(self.sd_model)
         crop_region = None
 
@@ -347,11 +438,14 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
             else:
                 self.image_mask = images.resize_image(self.resize_mode, self.image_mask, self.width, self.height)
                 np_mask = np.array(self.image_mask)
-                np_mask = 255 - np.clip((255 - np_mask.astype(np.float)) * 2, 0, 255).astype(np.uint8)
+                np_mask = np.clip((np_mask.astype(np.float32)) * 2, 0, 255).astype(np.uint8)
                 self.mask_for_overlay = Image.fromarray(np_mask)
 
             self.overlay_images = []
 
+        latent_mask = self.latent_mask if self.latent_mask is not None else self.image_mask
+
+        self.color_corrections = []
         imgs = []
         for img in self.init_images:
             image = img.convert("RGB")
@@ -360,9 +454,6 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
                 image = images.resize_image(self.resize_mode, image, self.width, self.height)
 
             if self.image_mask is not None:
-                if self.inpainting_fill != 1:
-                    image = fill(image, self.mask_for_overlay)
-
                 image_masked = Image.new('RGBa', (image.width, image.height))
                 image_masked.paste(image.convert("RGBA").convert("RGBa"), mask=ImageOps.invert(self.mask_for_overlay.convert('L')))
 
@@ -372,6 +463,13 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
                 image = image.crop(crop_region)
                 image = images.resize_image(2, image, self.width, self.height)
 
+            if self.image_mask is not None:
+                if self.inpainting_fill != 1:
+                    image = fill(image, latent_mask)
+
+            if opts.img2img_color_correction:
+                self.color_corrections.append(setup_color_correction(image))
+
             image = np.array(image).astype(np.float32) / 255.0
             image = np.moveaxis(image, 2, 0)
 
@@ -394,17 +492,18 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
         self.init_latent = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(image))
 
         if self.image_mask is not None:
-            init_mask = self.latent_mask if self.latent_mask is not None else self.image_mask
+            init_mask = latent_mask
             latmask = init_mask.convert('RGB').resize((self.init_latent.shape[3], self.init_latent.shape[2]))
-            latmask = np.moveaxis(np.array(latmask, dtype=np.float64), 2, 0) / 255
+            latmask = np.moveaxis(np.array(latmask, dtype=np.float32), 2, 0) / 255
             latmask = latmask[0]
+            latmask = np.around(latmask)
             latmask = np.tile(latmask[None], (4, 1, 1))
 
             self.mask = torch.asarray(1.0 - latmask).to(shared.device).type(self.sd_model.dtype)
             self.nmask = torch.asarray(latmask).to(shared.device).type(self.sd_model.dtype)
 
             if self.inpainting_fill == 2:
-                self.init_latent = self.init_latent * self.mask + create_random_tensors(self.init_latent.shape[1:], [self.seed + x + 1 for x in range(self.init_latent.shape[0])]) * self.nmask
+                self.init_latent = self.init_latent * self.mask + create_random_tensors(self.init_latent.shape[1:], [seed + x + 1 for x in range(self.init_latent.shape[0])]) * self.nmask
             elif self.inpainting_fill == 3:
                 self.init_latent = self.init_latent * self.mask
 

modules/prompt_parser.py

@@ -0,0 +1,128 @@
+import re
+from collections import namedtuple
+import torch
+
+import modules.shared as shared
+
+re_prompt = re.compile(r'''
+(.*?)
+\[
+    ([^]:]+):
+    (?:([^]:]*):)?
+    ([0-9]*\.?[0-9]+)
+]
+|
+(.+)
+''', re.X)
+
+# a prompt like this: "fantasy landscape with a [mountain:lake:0.25] and [an oak:a christmas tree:0.75][ in foreground::0.6][ in background:0.25] [shoddy:masterful:0.5]"
+# will be represented with prompt_schedule like this (assuming steps=100):
+# [25, 'fantasy landscape with a mountain and an oak in foreground shoddy']
+# [50, 'fantasy landscape with a lake and an oak in foreground in background shoddy']
+# [60, 'fantasy landscape with a lake and an oak in foreground in background masterful']
+# [75, 'fantasy landscape with a lake and an oak in background masterful']
+# [100, 'fantasy landscape with a lake and a christmas tree in background masterful']
+
+
+def get_learned_conditioning_prompt_schedules(prompts, steps):
+    res = []
+    cache = {}
+
+    for prompt in prompts:
+        prompt_schedule: list[list[str | int]] = [[steps, ""]]
+
+        cached = cache.get(prompt, None)
+        if cached is not None:
+            res.append(cached)
+
+        for m in re_prompt.finditer(prompt):
+            plaintext = m.group(1) if m.group(5) is None else m.group(5)
+            concept_from = m.group(2)
+            concept_to = m.group(3)
+            if concept_to is None:
+                concept_to = concept_from
+                concept_from = ""
+            swap_position = float(m.group(4)) if m.group(4) is not None else None
+
+            if swap_position is not None:
+                if swap_position < 1:
+                    swap_position = swap_position * steps
+                swap_position = int(min(swap_position, steps))
+
+            swap_index = None
+            found_exact_index = False
+            for i in range(len(prompt_schedule)):
+                end_step = prompt_schedule[i][0]
+                prompt_schedule[i][1] += plaintext
+
+                if swap_position is not None and swap_index is None:
+                    if swap_position == end_step:
+                        swap_index = i
+                        found_exact_index = True
+
+                    if swap_position < end_step:
+                        swap_index = i
+
+            if swap_index is not None:
+                if not found_exact_index:
+                    prompt_schedule.insert(swap_index, [swap_position, prompt_schedule[swap_index][1]])
+
+                for i in range(len(prompt_schedule)):
+                    end_step = prompt_schedule[i][0]
+                    must_replace = swap_position < end_step
+
+                    prompt_schedule[i][1] += concept_to if must_replace else concept_from
+
+        res.append(prompt_schedule)
+        cache[prompt] = prompt_schedule
+        #for t in prompt_schedule:
+        #    print(t)
+
+    return res
+
+
+ScheduledPromptConditioning = namedtuple("ScheduledPromptConditioning", ["end_at_step", "cond"])
+ScheduledPromptBatch = namedtuple("ScheduledPromptBatch", ["shape", "schedules"])
+
+
+def get_learned_conditioning(prompts, steps):
+
+    res = []
+
+    prompt_schedules = get_learned_conditioning_prompt_schedules(prompts, steps)
+    cache = {}
+
+    for prompt, prompt_schedule in zip(prompts, prompt_schedules):
+
+        cached = cache.get(prompt, None)
+        if cached is not None:
+            res.append(cached)
+
+        texts = [x[1] for x in prompt_schedule]
+        conds = shared.sd_model.get_learned_conditioning(texts)
+
+        cond_schedule = []
+        for i, (end_at_step, text) in enumerate(prompt_schedule):
+            cond_schedule.append(ScheduledPromptConditioning(end_at_step, conds[i]))
+
+        cache[prompt] = cond_schedule
+        res.append(cond_schedule)
+
+    return ScheduledPromptBatch((len(prompts),) + res[0][0].cond.shape, res)
+
+
+def reconstruct_cond_batch(c: ScheduledPromptBatch, current_step):
+    res = torch.zeros(c.shape)
+    for i, cond_schedule in enumerate(c.schedules):
+        target_index = 0
+        for curret_index, (end_at, cond) in enumerate(cond_schedule):
+            if current_step <= end_at:
+                target_index = curret_index
+                break
+        res[i] = cond_schedule[target_index].cond
+
+    return res.to(shared.device)
+
+
+
+#get_learned_conditioning_prompt_schedules(["fantasy landscape with a [mountain:lake:0.25] and [an oak:a christmas tree:0.75][ in foreground::0.6][ in background:0.25] [shoddy:masterful:0.5]"], 100)

modules/realesrgan_model.py

@@ -5,7 +5,7 @@ import numpy as np
 from PIL import Image
 
 import modules.images
-from modules.shared import cmd_opts
+from modules.shared import cmd_opts, opts
 
 RealesrganModelInfo = namedtuple("RealesrganModelInfo", ["name", "location", "model", "netscale"])
 
@@ -76,7 +76,9 @@ def upscale_with_realesrgan(image, RealESRGAN_upscaling, RealESRGAN_model_index)
         scale=info.netscale,
         model_path=info.location,
         model=model,
-        half=not cmd_opts.no_half
+        half=not cmd_opts.no_half,
+        tile=opts.ESRGAN_tile,
+        tile_pad=opts.ESRGAN_tile_overlap,
     )
 
     upsampled = upsampler.enhance(np.array(image), outscale=RealESRGAN_upscaling)[0]

modules/safety.py

@@ -0,0 +1,42 @@
+import torch
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+from transformers import AutoFeatureExtractor
+from PIL import Image
+
+import modules.shared as shared
+
+safety_model_id = "CompVis/stable-diffusion-safety-checker"
+safety_feature_extractor = None
+safety_checker = None
+
+def numpy_to_pil(images):
+    """
+    Convert a numpy image or a batch of images to a PIL image.
+    """
+    if images.ndim == 3:
+        images = images[None, ...]
+    images = (images * 255).round().astype("uint8")
+    pil_images = [Image.fromarray(image) for image in images]
+
+    return pil_images
+
+# check and replace nsfw content
+def check_safety(x_image):
+    global safety_feature_extractor, safety_checker
+
+    if safety_feature_extractor is None:
+        safety_feature_extractor = AutoFeatureExtractor.from_pretrained(safety_model_id)
+        safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id)
+
+    safety_checker_input = safety_feature_extractor(numpy_to_pil(x_image), return_tensors="pt")
+    x_checked_image, has_nsfw_concept = safety_checker(images=x_image, clip_input=safety_checker_input.pixel_values)
+
+    return x_checked_image, has_nsfw_concept
+
+
+def censor_batch(x):
+    x_samples_ddim_numpy = x.cpu().permute(0, 2, 3, 1).numpy()
+    x_checked_image, has_nsfw_concept = check_safety(x_samples_ddim_numpy)
+    x = torch.from_numpy(x_checked_image).permute(0, 3, 1, 2)
+
+    return x

modules/scripts.py

@@ -6,6 +6,7 @@ import modules.ui as ui
 import gradio as gr
 
 from modules.processing import StableDiffusionProcessing
+from modules import shared
 
 class Script:
     filename = None
@@ -41,10 +42,10 @@ def load_scripts(basedir):
         if not os.path.isfile(path):
             continue
 
-        with open(path, "r", encoding="utf8") as file:
-            text = file.read()
-
         try:
+            with open(path, "r", encoding="utf8") as file:
+                text = file.read()
+
             from types import ModuleType
             compiled = compile(text, path, 'exec')
             module = ModuleType(filename)
@@ -137,6 +138,8 @@ class ScriptRunner:
         script_args = args[script.args_from:script.args_to]
         processed = script.run(p, *script_args)
 
+        shared.total_tqdm.clear()
+
         return processed
 
 

modules/sd_hijack.py

@@ -1,3 +1,4 @@
+import math
 import os
 import sys
 import traceback
@@ -10,10 +11,11 @@ from modules.shared import opts, device, cmd_opts
 from ldm.util import default
 from einops import rearrange
 import ldm.modules.attention
+import ldm.modules.diffusionmodules.model
 
 
 # see https://github.com/basujindal/stable-diffusion/pull/117 for discussion
-def split_cross_attention_forward(self, x, context=None, mask=None):
+def split_cross_attention_forward_v1(self, x, context=None, mask=None):
     h = self.heads
 
     q = self.to_q(x)
@@ -42,6 +44,133 @@ def split_cross_attention_forward(self, x, context=None, mask=None):
     return self.to_out(r2)
 
 
+# taken from https://github.com/Doggettx/stable-diffusion
+def split_cross_attention_forward(self, x, context=None, mask=None):
+    h = self.heads
+
+    q_in = self.to_q(x)
+    context = default(context, x)
+    k_in = self.to_k(context)
+    v_in = self.to_v(context)
+    del context, x
+
+    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))
+    del q_in, k_in, v_in
+
+    r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device)
+
+    stats = torch.cuda.memory_stats(q.device)
+    mem_active = stats['active_bytes.all.current']
+    mem_reserved = stats['reserved_bytes.all.current']
+    mem_free_cuda, _ = torch.cuda.mem_get_info(torch.cuda.current_device())
+    mem_free_torch = mem_reserved - mem_active
+    mem_free_total = mem_free_cuda + mem_free_torch
+
+    gb = 1024 ** 3
+    tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * q.element_size()
+    modifier = 3 if q.element_size() == 2 else 2.5
+    mem_required = tensor_size * modifier
+    steps = 1
+
+    if mem_required > mem_free_total:
+        steps = 2 ** (math.ceil(math.log(mem_required / mem_free_total, 2)))
+        # print(f"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB "
+        #       f"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}")
+
+    if steps > 64:
+        max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64
+        raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). '
+                           f'Need: {mem_required / 64 / gb:0.1f}GB free, Have:{mem_free_total / gb:0.1f}GB free')
+
+    slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
+    for i in range(0, q.shape[1], slice_size):
+        end = i + slice_size
+        s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k) * self.scale
+
+        s2 = s1.softmax(dim=-1, dtype=q.dtype)
+        del s1
+
+        r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v)
+        del s2
+
+    del q, k, v
+
+    r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h)
+    del r1
+
+    return self.to_out(r2)
+
+def nonlinearity_hijack(x):
+    # swish
+    t = torch.sigmoid(x)
+    x *= t
+    del t
+
+    return x
+
+def cross_attention_attnblock_forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q1 = self.q(h_)
+        k1 = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b, c, h, w = q1.shape
+
+        q2 = q1.reshape(b, c, h*w)
+        del q1
+
+        q = q2.permute(0, 2, 1)   # b,hw,c
+        del q2
+
+        k = k1.reshape(b, c, h*w) # b,c,hw
+        del k1
+
+        h_ = torch.zeros_like(k, device=q.device)
+
+        stats = torch.cuda.memory_stats(q.device)
+        mem_active = stats['active_bytes.all.current']
+        mem_reserved = stats['reserved_bytes.all.current']
+        mem_free_cuda, _ = torch.cuda.mem_get_info(torch.cuda.current_device())
+        mem_free_torch = mem_reserved - mem_active
+        mem_free_total = mem_free_cuda + mem_free_torch
+
+        tensor_size = q.shape[0] * q.shape[1] * k.shape[2] * q.element_size()
+        mem_required = tensor_size * 2.5
+        steps = 1
+
+        if mem_required > mem_free_total:
+            steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2)))
+
+        slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
+        for i in range(0, q.shape[1], slice_size):
+            end = i + slice_size
+
+            w1 = torch.bmm(q[:, i:end], k)     # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+            w2 = w1 * (int(c)**(-0.5))
+            del w1
+            w3 = torch.nn.functional.softmax(w2, dim=2, dtype=q.dtype)
+            del w2
+
+            # attend to values
+            v1 = v.reshape(b, c, h*w)
+            w4 = w3.permute(0, 2, 1)   # b,hw,hw (first hw of k, second of q)
+            del w3
+
+            h_[:, :, i:end] = torch.bmm(v1, w4)     # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+            del v1, w4
+
+        h2 = h_.reshape(b, c, h, w)
+        del h_
+
+        h3 = self.proj_out(h2)
+        del h2
+
+        h3 += x
+
+        return h3
+
 class StableDiffusionModelHijack:
     ids_lookup = {}
     word_embeddings = {}
@@ -73,11 +202,21 @@ class StableDiffusionModelHijack:
             name = os.path.splitext(filename)[0]
 
             data = torch.load(path)
-            param_dict = data['string_to_param']
-            if hasattr(param_dict, '_parameters'):
-                param_dict = getattr(param_dict, '_parameters')  # fix for torch 1.12.1 loading saved file from torch 1.11
-            assert len(param_dict) == 1, 'embedding file has multiple terms in it'
-            emb = next(iter(param_dict.items()))[1]
+
+            # textual inversion embeddings
+            if 'string_to_param' in data:
+                param_dict = data['string_to_param']
+                if hasattr(param_dict, '_parameters'):
+                    param_dict = getattr(param_dict, '_parameters')  # fix for torch 1.12.1 loading saved file from torch 1.11
+                assert len(param_dict) == 1, 'embedding file has multiple terms in it'
+                emb = next(iter(param_dict.items()))[1]
+            elif type(data) == dict and type(next(iter(data.values()))) == torch.Tensor:
+                assert len(data.keys()) == 1, 'embedding file has multiple terms in it'
+
+                emb = next(iter(data.values()))
+                if len(emb.shape) == 1:
+                    emb = emb.unsqueeze(0)
+
             self.word_embeddings[name] = emb.detach()
             self.word_embeddings_checksums[name] = f'{const_hash(emb.reshape(-1))&0xffff:04x}'
 
@@ -106,6 +245,10 @@ class StableDiffusionModelHijack:
 
         if cmd_opts.opt_split_attention:
             ldm.modules.attention.CrossAttention.forward = split_cross_attention_forward
+            ldm.modules.diffusionmodules.model.nonlinearity = nonlinearity_hijack
+            ldm.modules.diffusionmodules.model.AttnBlock.forward = cross_attention_attnblock_forward
+        elif cmd_opts.opt_split_attention_v1:
+            ldm.modules.attention.CrossAttention.forward = split_cross_attention_forward_v1
 
         def flatten(el):
             flattened = [flatten(children) for children in el.children()]
@@ -232,7 +375,7 @@ class FrozenCLIPEmbedderWithCustomWords(torch.nn.Module):
         z = outputs.last_hidden_state
 
         # restoring original mean is likely not correct, but it seems to work well to prevent artifacts that happen otherwise
-        batch_multipliers = torch.asarray(np.array(batch_multipliers)).to(device)
+        batch_multipliers = torch.asarray(batch_multipliers).to(device)
         original_mean = z.mean()
         z *= batch_multipliers.reshape(batch_multipliers.shape + (1,)).expand(z.shape)
         new_mean = z.mean()

modules/sd_samplers.py

@@ -7,6 +7,7 @@ from PIL import Image
 import k_diffusion.sampling
 import ldm.models.diffusion.ddim
 import ldm.models.diffusion.plms
+from modules import prompt_parser
 
 from modules.shared import opts, cmd_opts, state
 import modules.shared as shared
@@ -53,27 +54,19 @@ def store_latent(decoded):
             shared.state.current_image = sample_to_image(decoded)
 
 
-def p_sample_ddim_hook(sampler_wrapper, x_dec, cond, ts, *args, **kwargs):
-    if sampler_wrapper.mask is not None:
-        img_orig = sampler_wrapper.sampler.model.q_sample(sampler_wrapper.init_latent, ts)
-        x_dec = img_orig * sampler_wrapper.mask + sampler_wrapper.nmask * x_dec
-
-    store_latent(x_dec)
-
-    return sampler_wrapper.orig_p_sample_ddim(x_dec, cond, ts, *args, **kwargs)
-
 
 def extended_tdqm(sequence, *args, desc=None, **kwargs):
     state.sampling_steps = len(sequence)
     state.sampling_step = 0
 
-    for x in tqdm.tqdm(sequence, *args, desc=state.job, **kwargs):
+    for x in tqdm.tqdm(sequence, *args, desc=state.job, file=shared.progress_print_out, **kwargs):
         if state.interrupted:
             break
 
         yield x
 
         state.sampling_step += 1
+        shared.total_tqdm.update()
 
 
 ldm.models.diffusion.ddim.tqdm = lambda *args, desc=None, **kwargs: extended_tdqm(*args, desc=desc, **kwargs)
@@ -83,15 +76,34 @@ ldm.models.diffusion.plms.tqdm = lambda *args, desc=None, **kwargs: extended_tdq
 class VanillaStableDiffusionSampler:
     def __init__(self, constructor, sd_model):
         self.sampler = constructor(sd_model)
-        self.orig_p_sample_ddim = self.sampler.p_sample_ddim if hasattr(self.sampler, 'p_sample_ddim') else None
+        self.orig_p_sample_ddim = self.sampler.p_sample_ddim if hasattr(self.sampler, 'p_sample_ddim') else self.sampler.p_sample_plms
         self.mask = None
         self.nmask = None
         self.init_latent = None
+        self.step = 0
+
+    def p_sample_ddim_hook(self, x_dec, cond, ts, unconditional_conditioning, *args, **kwargs):
+        cond = prompt_parser.reconstruct_cond_batch(cond, self.step)
+        unconditional_conditioning = prompt_parser.reconstruct_cond_batch(unconditional_conditioning, self.step)
+
+        if self.mask is not None:
+            img_orig = self.sampler.model.q_sample(self.init_latent, ts)
+            x_dec = img_orig * self.mask + self.nmask * x_dec
+
+        res = self.orig_p_sample_ddim(x_dec, cond, ts, unconditional_conditioning=unconditional_conditioning, *args, **kwargs)
+
+        if self.mask is not None:
+            store_latent(self.init_latent * self.mask + self.nmask * res[1])
+        else:
+            store_latent(res[1])
+
+        self.step += 1
+        return res
 
     def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning):
         t_enc = int(min(p.denoising_strength, 0.999) * p.steps)
 
-        # existing code fails with cetin step counts, like 9
+        # existing code fails with cetain step counts, like 9
         try:
             self.sampler.make_schedule(ddim_num_steps=p.steps, verbose=False)
         except Exception:
@@ -99,7 +111,7 @@ class VanillaStableDiffusionSampler:
 
         x1 = self.sampler.stochastic_encode(x, torch.tensor([t_enc] * int(x.shape[0])).to(shared.device), noise=noise)
 
-        self.sampler.p_sample_ddim = lambda x_dec, cond, ts, *args, **kwargs: p_sample_ddim_hook(self, x_dec, cond, ts, *args, **kwargs)
+        self.sampler.p_sample_ddim = self.p_sample_ddim_hook
         self.mask = p.mask
         self.nmask = p.nmask
         self.init_latent = p.init_latent
@@ -109,7 +121,19 @@ class VanillaStableDiffusionSampler:
         return samples
 
     def sample(self, p, x, conditioning, unconditional_conditioning):
-        samples_ddim, _ = self.sampler.sample(S=p.steps, conditioning=conditioning, batch_size=int(x.shape[0]), shape=x[0].shape, verbose=False, unconditional_guidance_scale=p.cfg_scale, unconditional_conditioning=unconditional_conditioning, x_T=x)
+        for fieldname in ['p_sample_ddim', 'p_sample_plms']:
+            if hasattr(self.sampler, fieldname):
+                setattr(self.sampler, fieldname, self.p_sample_ddim_hook)
+        self.mask = None
+        self.nmask = None
+        self.init_latent = None
+
+        # existing code fails with cetin step counts, like 9
+        try:
+            samples_ddim, _ = self.sampler.sample(S=p.steps, conditioning=conditioning, batch_size=int(x.shape[0]), shape=x[0].shape, verbose=False, unconditional_guidance_scale=p.cfg_scale, unconditional_conditioning=unconditional_conditioning, x_T=x)
+        except Exception:
+            samples_ddim, _ = self.sampler.sample(S=p.steps+1, conditioning=conditioning, batch_size=int(x.shape[0]), shape=x[0].shape, verbose=False, unconditional_guidance_scale=p.cfg_scale, unconditional_conditioning=unconditional_conditioning, x_T=x)
+
         return samples_ddim
 
 
@@ -120,8 +144,12 @@ class CFGDenoiser(torch.nn.Module):
         self.mask = None
         self.nmask = None
         self.init_latent = None
+        self.step = 0
 
     def forward(self, x, sigma, uncond, cond, cond_scale):
+        cond = prompt_parser.reconstruct_cond_batch(cond, self.step)
+        uncond = prompt_parser.reconstruct_cond_batch(uncond, self.step)
+
         if shared.batch_cond_uncond:
             x_in = torch.cat([x] * 2)
             sigma_in = torch.cat([sigma] * 2)
@@ -136,6 +164,8 @@ class CFGDenoiser(torch.nn.Module):
         if self.mask is not None:
             denoised = self.init_latent * self.mask + self.nmask * denoised
 
+        self.step += 1
+
         return denoised
 
 
@@ -143,13 +173,14 @@ def extended_trange(count, *args, **kwargs):
     state.sampling_steps = count
     state.sampling_step = 0
 
-    for x in tqdm.trange(count, *args, desc=state.job, **kwargs):
+    for x in tqdm.trange(count, *args, desc=state.job, file=shared.progress_print_out, **kwargs):
         if state.interrupted:
             break
 
         yield x
 
         state.sampling_step += 1
+        shared.total_tqdm.update()
 
 
 class KDiffusionSampler:
@@ -165,6 +196,7 @@ class KDiffusionSampler:
     def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning):
         t_enc = int(min(p.denoising_strength, 0.999) * p.steps)
         sigmas = self.model_wrap.get_sigmas(p.steps)
+
         noise = noise * sigmas[p.steps - t_enc - 1]
 
         xi = x + noise

modules/shared.py

@@ -1,15 +1,17 @@
+import sys
 import argparse
 import json
 import os
 
 import gradio as gr
 import torch
+import tqdm
 
 import modules.artists
 from modules.paths import script_path, sd_path
-import modules.codeformer_model
-
-config_filename = "config.json"
+from modules.devices import get_optimal_device
+import modules.styles
+import modules.interrogate
 
 sd_model_file = os.path.join(script_path, 'model.ckpt')
 if not os.path.exists(sd_model_file):
@@ -21,27 +23,38 @@ parser.add_argument("--ckpt", type=str, default=os.path.join(sd_path, sd_model_f
 parser.add_argument("--gfpgan-dir", type=str, help="GFPGAN directory", default=('./src/gfpgan' if os.path.exists('./src/gfpgan') else './GFPGAN'))
 parser.add_argument("--gfpgan-model", type=str, help="GFPGAN model file name", default='GFPGANv1.3.pth')
 parser.add_argument("--no-half", action='store_true', help="do not switch the model to 16-bit floats")
-parser.add_argument("--no-progressbar-hiding", action='store_true', help="do not hide progressbar in gradio UI (we hide it because it slows down ML if you have hardware accleration in browser)")
+parser.add_argument("--no-progressbar-hiding", action='store_true', help="do not hide progressbar in gradio UI (we hide it because it slows down ML if you have hardware acceleration in browser)")
 parser.add_argument("--max-batch-count", type=int, default=16, help="maximum batch count value for the UI")
-parser.add_argument("--embeddings-dir", type=str, default='embeddings', help="embeddings dirtectory for textual inversion (default: embeddings)")
+parser.add_argument("--embeddings-dir", type=str, default=os.path.join(script_path, 'embeddings'), help="embeddings directory for textual inversion (default: embeddings)")
 parser.add_argument("--allow-code", action='store_true', help="allow custom script execution from webui")
-parser.add_argument("--medvram", action='store_true', help="enable stable diffusion model optimizations for sacrficing a little speed for low VRM usage")
-parser.add_argument("--lowvram", action='store_true', help="enable stable diffusion model optimizations for sacrficing a lot of speed for very low VRM usage")
-parser.add_argument("--always-batch-cond-uncond", action='store_true', help="a workaround test; may help with speed in you use --lowvram")
-parser.add_argument("--unload-gfpgan", action='store_true', help="unload GFPGAN every time after processing images. Warning: seems to cause memory leaks")
+parser.add_argument("--medvram", action='store_true', help="enable stable diffusion model optimizations for sacrificing a little speed for low VRM usage")
+parser.add_argument("--lowvram", action='store_true', help="enable stable diffusion model optimizations for sacrificing a lot of speed for very low VRM usage")
+parser.add_argument("--always-batch-cond-uncond", action='store_true', help="a workaround test; may help with speed if you use --lowvram")
+parser.add_argument("--unload-gfpgan", action='store_true', help="does not do anything.")
 parser.add_argument("--precision", type=str, help="evaluate at this precision", choices=["full", "autocast"], default="autocast")
 parser.add_argument("--share", action='store_true', help="use share=True for gradio and make the UI accessible through their site (doesn't work for me but you might have better luck)")
 parser.add_argument("--esrgan-models-path", type=str, help="path to directory with ESRGAN models", default=os.path.join(script_path, 'ESRGAN'))
-parser.add_argument("--opt-split-attention", action='store_true', help="enable optimization that reduced vram usage by a lot for about 10%% decrease in performance")
+parser.add_argument("--opt-split-attention", action='store_true', help="enable optimization that reduce vram usage by a lot for about 10%% decrease in performance")
+parser.add_argument("--opt-split-attention-v1", action='store_true', help="enable older version of --opt-split-attention optimization")
 parser.add_argument("--listen", action='store_true', help="launch gradio with 0.0.0.0 as server name, allowing to respond to network requests")
+parser.add_argument("--port", type=int, help="launch gradio with given server port, you need root/admin rights for ports < 1024, defaults to 7860 if available", default=None)
+parser.add_argument("--show-negative-prompt", action='store_true', help="does not do anything", default=False)
+parser.add_argument("--ui-config-file", type=str, help="filename to use for ui configuration", default=os.path.join(script_path, 'ui-config.json'))
+parser.add_argument("--hide-ui-dir-config", action='store_true', help="hide directory configuration from webui", default=False)
+parser.add_argument("--ui-settings-file", type=str, help="filename to use for ui settings", default=os.path.join(script_path, 'config.json'))
+parser.add_argument("--gradio-debug",  action='store_true', help="launch gradio with --debug option")
+parser.add_argument("--gradio-auth", type=str, help='set gradio authentication like "username:password"; or comma-delimit multiple like "u1:p1,u2:p2,u3:p3"', default=None)
+parser.add_argument("--opt-channelslast", action='store_true', help="change memory type for stable diffusion to channels last")
+parser.add_argument("--styles-file", type=str, help="filename to use for styles", default=os.path.join(script_path, 'styles.csv'))
+
 cmd_opts = parser.parse_args()
 
-cpu = torch.device("cpu")
-gpu = torch.device("cuda")
-device = gpu if torch.cuda.is_available() else cpu
+device = get_optimal_device()
+
 batch_cond_uncond = cmd_opts.always_batch_cond_uncond or not (cmd_opts.lowvram or cmd_opts.medvram)
 parallel_processing_allowed = not cmd_opts.lowvram and not cmd_opts.medvram
 
+config_filename = cmd_opts.ui_settings_file
 
 class State:
     interrupted = False
@@ -54,7 +67,6 @@ class State:
     current_image = None
     current_image_sampling_step = 0
 
-
     def interrupt(self):
         self.interrupted = True
 
@@ -68,18 +80,13 @@ state = State()
 
 artist_db = modules.artists.ArtistsDatabase(os.path.join(script_path, 'artists.csv'))
 
+styles_filename = cmd_opts.styles_file
+prompt_styles = modules.styles.StyleDatabase(styles_filename)
+
+interrogator = modules.interrogate.InterrogateModels("interrogate")
+
 face_restorers = []
 
-def find_any_font():
-    fonts = ['/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf']
-
-    for font in fonts:
-        if os.path.exists(font):
-            return font
-
-    return "arial.ttf"
-
-
 class Options:
     class OptionInfo:
         def __init__(self, default=None, label="", component=None, component_args=None):
@@ -89,20 +96,25 @@ class Options:
             self.component_args = component_args
 
     data = None
+    hide_dirs = {"visible": False} if cmd_opts.hide_ui_dir_config else None
     data_labels = {
-        "outdir_samples": OptionInfo("", "Output dictectory for images; if empty, defaults to two directories below"),
-        "outdir_txt2img_samples": OptionInfo("outputs/txt2img-images", 'Output dictectory for txt2img images'),
-        "outdir_img2img_samples": OptionInfo("outputs/img2img-images", 'Output dictectory for img2img images'),
-        "outdir_extras_samples": OptionInfo("outputs/extras-images", 'Output dictectory for images from extras tab'),
-        "outdir_grids": OptionInfo("", "Output dictectory for grids; if empty, defaults to two directories below"),
-        "outdir_txt2img_grids": OptionInfo("outputs/txt2img-grids", 'Output dictectory for txt2img grids'),
-        "outdir_img2img_grids": OptionInfo("outputs/img2img-grids", 'Output dictectory for img2img grids'),
-        "save_to_dirs": OptionInfo(False, "When writing images/grids, create a directory with name derived from the prompt"),
-        "save_to_dirs_prompt_len": OptionInfo(10, "When using above, how many words from prompt to put into directory name", gr.Slider, {"minimum": 1, "maximum": 32, "step": 1}),
-        "outdir_save": OptionInfo("log/images", "Directory for saving images using the Save button"),
-        "samples_save": OptionInfo(True, "Save indiviual samples"),
-        "samples_format": OptionInfo('png', 'File format for indiviual samples'),
-        "grid_save": OptionInfo(True, "Save image grids"),
+        "samples_filename_pattern": OptionInfo("", "Images filename pattern"),
+        "save_to_dirs": OptionInfo(False, "Save images to a subdirectory"),
+        "grid_save_to_dirs": OptionInfo(False, "Save grids to subdirectory"),
+        "directories_filename_pattern": OptionInfo("", "Directory name pattern"),
+        "outdir_samples": OptionInfo("", "Output directory for images; if empty, defaults to two directories below", component_args=hide_dirs),
+        "outdir_txt2img_samples": OptionInfo("outputs/txt2img-images", 'Output directory for txt2img images', component_args=hide_dirs),
+        "outdir_img2img_samples": OptionInfo("outputs/img2img-images", 'Output directory for img2img images', component_args=hide_dirs),
+        "outdir_extras_samples": OptionInfo("outputs/extras-images", 'Output directory for images from extras tab', component_args=hide_dirs),
+        "outdir_grids": OptionInfo("", "Output directory for grids; if empty, defaults to two directories below", component_args=hide_dirs),
+        "outdir_txt2img_grids": OptionInfo("outputs/txt2img-grids", 'Output directory for txt2img grids', component_args=hide_dirs),
+        "outdir_img2img_grids": OptionInfo("outputs/img2img-grids", 'Output directory for img2img grids', component_args=hide_dirs),
+        "outdir_save": OptionInfo("log/images", "Directory for saving images using the Save button", component_args=hide_dirs),
+        "samples_save": OptionInfo(True, "Always save all generated images"),
+        "save_selected_only": OptionInfo(False, "When using 'Save' button, only save a single selected image"),
+        "samples_format": OptionInfo('png', 'File format for individual samples'),
+        "filter_nsfw": OptionInfo(False, "Filter NSFW content"),
+        "grid_save": OptionInfo(True, "Always save all generated image grids"),
         "return_grid": OptionInfo(True, "Show grid in results for web"),
         "grid_format": OptionInfo('png', 'File format for grids'),
         "grid_extended_filename": OptionInfo(False, "Add extended info (seed, prompt) to filename when saving grid"),
@@ -111,17 +123,28 @@ class Options:
         "jpeg_quality": OptionInfo(80, "Quality for saved jpeg images", gr.Slider, {"minimum": 1, "maximum": 100, "step": 1}),
         "export_for_4chan": OptionInfo(True, "If PNG image is larger than 4MB or any dimension is larger than 4000, downscale and save copy as JPG"),
         "enable_pnginfo": OptionInfo(True, "Save text information about generation parameters as chunks to png files"),
-        "font": OptionInfo(find_any_font(), "Font for image grids  that have text"),
-        "enable_emphasis": OptionInfo(True, "Use (text) to make model pay more attention to text text and [text] to make it pay less attention"),
+        "add_model_hash_to_info": OptionInfo(False, "Add model hash to generation information"),
+        "img2img_color_correction": OptionInfo(False, "Apply color correction to img2img results to match original colors."),
+        "font": OptionInfo("", "Font for image grids that have text"),
+        "enable_emphasis": OptionInfo(True, "Use (text) to make model pay more attention to text and [text] to make it pay less attention"),
         "save_txt": OptionInfo(False, "Create a text file next to every image with generation parameters."),
-        "ESRGAN_tile": OptionInfo(192, "Tile size for ESRGAN upscaling. 0 = no tiling.", gr.Slider, {"minimum": 0, "maximum": 512, "step": 16}),
-        "ESRGAN_tile_overlap": OptionInfo(8, "Tile overlap, in pixels for ESRGAN upscaling. Low values = visible seam.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}),
+        "ESRGAN_tile": OptionInfo(192, "Tile size for upscaling. 0 = no tiling.", gr.Slider, {"minimum": 0, "maximum": 512, "step": 16}),
+        "ESRGAN_tile_overlap": OptionInfo(8, "Tile overlap, in pixels for upscaling. Low values = visible seam.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}),
         "random_artist_categories": OptionInfo([], "Allowed categories for random artists selection when using the Roll button", gr.CheckboxGroup, {"choices": artist_db.categories()}),
         "upscale_at_full_resolution_padding": OptionInfo(16, "Inpainting at full resolution: padding, in pixels, for the masked region.", gr.Slider, {"minimum": 0, "maximum": 128, "step": 4}),
         "show_progressbar": OptionInfo(True, "Show progressbar"),
         "show_progress_every_n_steps": OptionInfo(0, "Show show image creation progress every N sampling steps. Set 0 to disable.", gr.Slider, {"minimum": 0, "maximum": 32, "step": 1}),
+        "multiple_tqdm": OptionInfo(True, "Add a second progress bar to the console that shows progress for an entire job. Broken in PyCharm console."),
         "face_restoration_model": OptionInfo(None, "Face restoration model", gr.Radio, lambda: {"choices": [x.name() for x in face_restorers]}),
         "code_former_weight": OptionInfo(0.5, "CodeFormer weight parameter; 0 = maximum effect; 1 = minimum effect", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01}),
+        "save_images_before_face_restoration": OptionInfo(False, "Save a copy of image before doing face restoration."),
+        "face_restoration_unload": OptionInfo(False, "Move face restoration model from VRAM into RAM after processing"),
+        "interrogate_keep_models_in_memory": OptionInfo(False, "Interrogate: keep models in VRAM"),
+        "interrogate_use_builtin_artists": OptionInfo(True, "Interrogate: use artists from artists.csv"),
+        "interrogate_clip_num_beams": OptionInfo(1, "Interrogate: num_beams for BLIP", gr.Slider, {"minimum": 1, "maximum": 16, "step": 1}),
+        "interrogate_clip_min_length": OptionInfo(24, "Interrogate: minimum description length (excluding artists, etc..)", gr.Slider, {"minimum": 1, "maximum": 128, "step": 1}),
+        "interrogate_clip_max_length": OptionInfo(48, "Interrogate: maximum description length", gr.Slider, {"minimum": 1, "maximum": 256, "step": 1}),
+        "interrogate_clip_dict_limit": OptionInfo(1500, "Interrogate: maximum number of lines in text file (0 = No limit)"),
     }
 
     def __init__(self):
@@ -160,5 +183,34 @@ if os.path.exists(config_filename):
 sd_upscalers = []
 
 sd_model = None
+sd_model_hash = ''
+
+progress_print_out = sys.stdout
 
 
+class TotalTQDM:
+    def __init__(self):
+        self._tqdm = None
+
+    def reset(self):
+        self._tqdm = tqdm.tqdm(
+            desc="Total progress",
+            total=state.job_count * state.sampling_steps,
+            position=1,
+            file=progress_print_out
+        )
+
+    def update(self):
+        if not opts.multiple_tqdm:
+            return
+        if self._tqdm is None:
+            self.reset()
+        self._tqdm.update()
+
+    def clear(self):
+        if self._tqdm is not None:
+            self._tqdm.close()
+            self._tqdm = None
+
+
+total_tqdm = TotalTQDM()

modules/styles.py

@@ -0,0 +1,86 @@
+# We need this so Python doesn't complain about the unknown StableDiffusionProcessing-typehint at runtime
+from __future__ import annotations
+
+import csv
+import os
+import os.path
+import typing
+import collections.abc as abc
+import tempfile
+import shutil
+
+if typing.TYPE_CHECKING:
+    # Only import this when code is being type-checked, it doesn't have any effect at runtime
+    from .processing import StableDiffusionProcessing
+
+
+class PromptStyle(typing.NamedTuple):
+    name: str
+    prompt: str
+    negative_prompt: str
+
+
+def merge_prompts(style_prompt: str, prompt: str) -> str:
+    if "{prompt}" in style_prompt:
+        res = style_prompt.replace("{prompt}", prompt)
+    else:
+        parts = filter(None, (prompt.strip(), style_prompt.strip()))
+        res = ", ".join(parts)
+
+    return res
+
+
+def apply_styles_to_prompt(prompt, styles):
+    for style in styles:
+        prompt = merge_prompts(style, prompt)
+
+    return prompt
+
+
+class StyleDatabase:
+    def __init__(self, path: str):
+        self.no_style = PromptStyle("None", "", "")
+        self.styles = {"None": self.no_style}
+
+        if not os.path.exists(path):
+            return
+
+        with open(path, "r", encoding="utf8", newline='') as file:
+            reader = csv.DictReader(file)
+            for row in reader:
+                # Support loading old CSV format with "name, text"-columns
+                prompt = row["prompt"] if "prompt" in row else row["text"]
+                negative_prompt = row.get("negative_prompt", "")
+                self.styles[row["name"]] = PromptStyle(row["name"], prompt, negative_prompt)
+
+    def apply_styles_to_prompt(self, prompt, styles):
+        return apply_styles_to_prompt(prompt, [self.styles.get(x, self.no_style).prompt for x in styles])
+
+    def apply_negative_styles_to_prompt(self, prompt, styles):
+        return apply_styles_to_prompt(prompt, [self.styles.get(x, self.no_style).negative_prompt for x in styles])
+
+    def apply_styles(self, p: StableDiffusionProcessing) -> None:
+        if isinstance(p.prompt, list):
+            p.prompt = [self.apply_styles_to_prompt(prompt, p.styles) for prompt in p.prompt]
+        else:
+            p.prompt = self.apply_styles_to_prompt(p.prompt, p.styles)
+
+        if isinstance(p.negative_prompt, list):
+            p.negative_prompt = [self.apply_negative_styles_to_prompt(prompt, p.styles) for prompt in p.negative_prompt]
+        else:
+            p.negative_prompt = self.apply_negative_styles_to_prompt(p.negative_prompt, p.styles)
+
+    def save_styles(self, path: str) -> None:
+        # Write to temporary file first, so we don't nuke the file if something goes wrong
+        fd, temp_path = tempfile.mkstemp(".csv")
+        with os.fdopen(fd, "w", encoding="utf8", newline='') as file:
+            # _fields is actually part of the public API: typing.NamedTuple is a replacement for collections.NamedTuple,
+            # and collections.NamedTuple has explicit documentation for accessing _fields. Same goes for _asdict()
+            writer = csv.DictWriter(file, fieldnames=PromptStyle._fields)
+            writer.writeheader()
+            writer.writerows(style._asdict() for k,     style in self.styles.items())
+
+        # Always keep a backup file around
+        if os.path.exists(path):
+            shutil.move(path, path + ".bak")
+        shutil.move(temp_path, path)

modules/txt2img.py

@@ -6,14 +6,19 @@ import modules.processing as processing
 from modules.ui import plaintext_to_html
 
 
-def txt2img(prompt: str, negative_prompt: str, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, height: int, width: int, *args):
+def txt2img(prompt: str, negative_prompt: str, prompt_style: str, prompt_style2: str, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, height: int, width: int, *args):
     p = StableDiffusionProcessingTxt2Img(
         sd_model=shared.sd_model,
         outpath_samples=opts.outdir_samples or opts.outdir_txt2img_samples,
         outpath_grids=opts.outdir_grids or opts.outdir_txt2img_grids,
         prompt=prompt,
+        styles=[prompt_style, prompt_style2],
         negative_prompt=negative_prompt,
         seed=seed,
+        subseed=subseed,
+        subseed_strength=subseed_strength,
+        seed_resize_from_h=seed_resize_from_h,
+        seed_resize_from_w=seed_resize_from_w,
         sampler_index=sampler_index,
         batch_size=batch_size,
         n_iter=n_iter,
@@ -25,6 +30,7 @@ def txt2img(prompt: str, negative_prompt: str, steps: int, sampler_index: int, r
         tiling=tiling,
     )
 
+    print(f"\ntxt2img: {prompt}", file=shared.progress_print_out)
     processed = modules.scripts.scripts_txt2img.run(p, *args)
 
     if processed is not None:
@@ -32,5 +38,7 @@ def txt2img(prompt: str, negative_prompt: str, steps: int, sampler_index: int, r
     else:
         processed = process_images(p)
 
+    shared.total_tqdm.clear()
+
     return processed.images, processed.js(), plaintext_to_html(processed.info)
 

modules/ui.py

@@ -2,6 +2,7 @@ import base64
 import html
 import io
 import json
+import math
 import mimetypes
 import os
 import random
@@ -25,6 +26,7 @@ import modules.realesrgan_model as realesrgan
 import modules.scripts
 import modules.gfpgan_model
 import modules.codeformer_model
+import modules.styles
 
 # this is a fix for Windows users. Without it, javascript files will be served with text/html content-type and the bowser will not show any UI
 mimetypes.init()
@@ -52,7 +54,7 @@ css_hide_progressbar = """
 """
 
 def plaintext_to_html(text):
-    text = "".join([f"<p>{html.escape(x)}</p>\n" for x in text.split('\n')])
+    text = "<p>" + "<br>\n".join([f"{html.escape(x)}" for x in text.split('\n')]) + "</p>"
     return text
 
 
@@ -78,7 +80,7 @@ def send_gradio_gallery_to_image(x):
     return image_from_url_text(x[0])
 
 
-def save_files(js_data, images):
+def save_files(js_data, images, index):
     import csv
 
     os.makedirs(opts.outdir_save, exist_ok=True)
@@ -86,12 +88,16 @@ def save_files(js_data, images):
     filenames = []
 
     data = json.loads(js_data)
+    
+    if index > -1 and opts.save_selected_only and (index > 0 or not opts.return_grid): # ensures we are looking at a specific non-grid picture, and we have save_selected_only
+        images = [images[index]]
+        data["seed"] += (index - 1 if opts.return_grid else index)
 
     with open(os.path.join(opts.outdir_save, "log.csv"), "a", encoding="utf8", newline='') as file:
         at_start = file.tell() == 0
         writer = csv.writer(file)
         if at_start:
-            writer.writerow(["prompt", "seed", "width", "height", "sampler", "cfgs", "steps", "filename"])
+            writer.writerow(["prompt", "seed", "width", "height", "sampler", "cfgs", "steps", "filename", "negative_prompt"])
 
         filename_base = str(int(time.time() * 1000))
         for i, filedata in enumerate(images):
@@ -106,7 +112,7 @@ def save_files(js_data, images):
 
             filenames.append(filename)
 
-        writer.writerow([data["prompt"], data["seed"], data["width"], data["height"], data["sampler"], data["cfg_scale"], data["steps"], filenames[0]])
+        writer.writerow([data["prompt"], data["seed"], data["width"], data["height"], data["sampler"], data["cfg_scale"], data["steps"], filenames[0], data["negative_prompt"]])
 
     return '', '', plaintext_to_html(f"Saved: {filenames[0]}")
 
@@ -192,14 +198,106 @@ def visit(x, func, path=""):
         func(path + "/" + str(x.label), x)
 
 
+def create_seed_inputs():
+    with gr.Row():
+        seed = gr.Number(label='Seed', value=-1)
+        subseed = gr.Number(label='Variation seed', value=-1, visible=False)
+        seed_checkbox = gr.Checkbox(label="Extra", elem_id="subseed_show", value=False)
+
+    with gr.Row():
+        subseed_strength = gr.Slider(label='Variation strength', value=0.0, minimum=0, maximum=1, step=0.01, visible=False)
+        seed_resize_from_w = gr.Slider(minimum=0, maximum=2048, step=64, label="Resize seed from width", value=0, visible=False)
+        seed_resize_from_h = gr.Slider(minimum=0, maximum=2048, step=64, label="Resize seed from height", value=0, visible=False)
+
+    def change_visiblity(show):
+
+        return {
+            subseed: gr_show(show),
+            subseed_strength: gr_show(show),
+            seed_resize_from_h: gr_show(show),
+            seed_resize_from_w: gr_show(show),
+        }
+
+    seed_checkbox.change(
+        change_visiblity,
+        inputs=[seed_checkbox],
+        outputs=[
+            subseed,
+            subseed_strength,
+            seed_resize_from_h,
+            seed_resize_from_w
+        ]
+    )
+
+    return seed, subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w
+
+
+def add_style(name: str, prompt: str, negative_prompt: str):
+    if name is None:
+        return [gr_show(), gr_show()]
+
+    style = modules.styles.PromptStyle(name, prompt, negative_prompt)
+    shared.prompt_styles.styles[style.name] = style
+    # Save all loaded prompt styles: this allows us to update the storage format in the future more easily, because we
+    # reserialize all styles every time we save them
+    shared.prompt_styles.save_styles(shared.styles_filename)
+
+    update = {"visible": True, "choices": list(shared.prompt_styles.styles), "__type__": "update"}
+    return [update, update, update, update]
+
+
+def apply_styles(prompt, prompt_neg, style1_name, style2_name):
+    prompt = shared.prompt_styles.apply_styles_to_prompt(prompt, [style1_name, style2_name])
+    prompt_neg = shared.prompt_styles.apply_negative_styles_to_prompt(prompt_neg, [style1_name, style2_name])
+
+    return [gr.Textbox.update(value=prompt), gr.Textbox.update(value=prompt_neg), gr.Dropdown.update(value="None"), gr.Dropdown.update(value="None")]
+
+
+def interrogate(image):
+    prompt = shared.interrogator.interrogate(image)
+
+    return gr_show(True) if prompt is None else prompt
+
+
+def create_toprow(is_img2img):
+    with gr.Row(elem_id="toprow"):
+        with gr.Column(scale=4):
+            with gr.Row():
+                with gr.Column(scale=8):
+                    with gr.Row():
+                        prompt = gr.Textbox(label="Prompt", elem_id="prompt", show_label=False, placeholder="Prompt", lines=2)
+                        roll = gr.Button('Roll', elem_id="roll", visible=len(shared.artist_db.artists) > 0)
+
+                with gr.Column(scale=1, elem_id="style_pos_col"):
+                    prompt_style = gr.Dropdown(label="Style 1", elem_id="style_index", choices=[k for k, v in shared.prompt_styles.styles.items()], value=next(iter(shared.prompt_styles.styles.keys())), visible=len(shared.prompt_styles.styles) > 1)
+
+            with gr.Row():
+                with gr.Column(scale=8):
+                    negative_prompt = gr.Textbox(label="Negative prompt", elem_id="negative_prompt", show_label=False, placeholder="Negative prompt", lines=2)
+
+                with gr.Column(scale=1, elem_id="style_neg_col"):
+                    prompt_style2 = gr.Dropdown(label="Style 2", elem_id="style2_index", choices=[k for k, v in shared.prompt_styles.styles.items()], value=next(iter(shared.prompt_styles.styles.keys())), visible=len(shared.prompt_styles.styles) > 1)
+
+        with gr.Column(scale=1):
+            with gr.Row():
+                submit = gr.Button('Generate', elem_id="generate", variant='primary')
+
+            with gr.Row():
+                if is_img2img:
+                    interrogate = gr.Button('Interrogate', elem_id="interrogate")
+                else:
+                    interrogate = None
+                prompt_style_apply = gr.Button('Apply style', elem_id="style_apply")
+                save_style = gr.Button('Create style', elem_id="style_create")
+
+            check_progress = gr.Button('Check progress', elem_id="check_progress", visible=False)
+
+    return prompt, roll, prompt_style, negative_prompt, prompt_style2, submit, interrogate, prompt_style_apply, save_style, check_progress
+
+
 def create_ui(txt2img, img2img, run_extras, run_pnginfo):
     with gr.Blocks(analytics_enabled=False) as txt2img_interface:
-        with gr.Row():
-            prompt = gr.Textbox(label="Prompt", elem_id="txt2img_prompt", show_label=False, placeholder="Prompt", lines=1)
-            negative_prompt = gr.Textbox(label="Negative prompt", elem_id="txt2img_negative_prompt", show_label=False, placeholder="Negative prompt", lines=1, visible=False)
-            roll = gr.Button('Roll', elem_id="txt2img_roll", visible=len(shared.artist_db.artists) > 0)
-            submit = gr.Button('Generate', elem_id="txt2img_generate", variant='primary')
-            check_progress = gr.Button('Check progress', elem_id="check_progress", visible=False)
+        txt2img_prompt, roll, txt2img_prompt_style, txt2img_negative_prompt, txt2img_prompt_style2, submit, _, txt2img_prompt_style_apply, txt2img_save_style, check_progress = create_toprow(is_img2img=False)
 
         with gr.Row().style(equal_height=False):
             with gr.Column(variant='panel'):
@@ -214,13 +312,13 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                     batch_count = gr.Slider(minimum=1, maximum=cmd_opts.max_batch_count, step=1, label='Batch count', value=1)
                     batch_size = gr.Slider(minimum=1, maximum=8, step=1, label='Batch size', value=1)
 
-                cfg_scale = gr.Slider(minimum=1.0, maximum=15.0, step=0.5, label='CFG Scale', value=7.0)
+                cfg_scale = gr.Slider(minimum=1.0, maximum=30.0, step=0.5, label='CFG Scale', value=7.0)
 
                 with gr.Group():
-                    height = gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512)
                     width = gr.Slider(minimum=64, maximum=2048, step=64, label="Width", value=512)
+                    height = gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512)
 
-                seed = gr.Number(label='Seed', value=-1)
+                seed, subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w = create_seed_inputs()
 
                 with gr.Group():
                     custom_inputs = modules.scripts.scripts_txt2img.setup_ui(is_img2img=False)
@@ -228,8 +326,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
             with gr.Column(variant='panel'):
                 with gr.Group():
                     txt2img_preview = gr.Image(elem_id='txt2img_preview', visible=False)
-                    txt2img_gallery = gr.Gallery(label='Output', elem_id='txt2img_gallery')
-
+                    txt2img_gallery = gr.Gallery(label='Output', elem_id='txt2img_gallery').style(grid=4)
 
                 with gr.Group():
                     with gr.Row():
@@ -245,13 +342,14 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                     html_info = gr.HTML()
                     generation_info = gr.Textbox(visible=False)
 
-
             txt2img_args = dict(
                 fn=txt2img,
                 _js="submit",
                 inputs=[
-                    prompt,
-                    negative_prompt,
+                    txt2img_prompt,
+                    txt2img_negative_prompt,
+                    txt2img_prompt_style,
+                    txt2img_prompt_style2,
                     steps,
                     sampler_index,
                     restore_faces,
@@ -260,6 +358,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                     batch_size,
                     cfg_scale,
                     seed,
+                    subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w,
                     height,
                     width,
                 ] + custom_inputs,
@@ -270,7 +369,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                 ]
             )
 
-            prompt.submit(**txt2img_args)
+            txt2img_prompt.submit(**txt2img_args)
             submit.click(**txt2img_args)
 
             check_progress.click(
@@ -280,7 +379,6 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                 outputs=[progressbar, txt2img_preview, txt2img_preview],
             )
 
-
             interrupt.click(
                 fn=lambda: shared.state.interrupt(),
                 inputs=[],
@@ -289,9 +387,11 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
 
             save.click(
                 fn=wrap_gradio_call(save_files),
+                _js = "(x, y, z) => [x, y, selected_gallery_index()]",
                 inputs=[
                     generation_info,
                     txt2img_gallery,
+                    html_info
                 ],
                 outputs=[
                     html_info,
@@ -303,28 +403,28 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
             roll.click(
                 fn=roll_artist,
                 inputs=[
-                    prompt,
+                    txt2img_prompt,
                 ],
                 outputs=[
-                    prompt
+                    txt2img_prompt,
                 ]
             )
 
-
     with gr.Blocks(analytics_enabled=False) as img2img_interface:
-        with gr.Row():
-            prompt = gr.Textbox(label="Prompt", elem_id="img2img_prompt", show_label=False, placeholder="Prompt", lines=1)
-            submit = gr.Button('Generate', elem_id="img2img_generate", variant='primary')
-            check_progress = gr.Button('Check progress', elem_id="check_progress", visible=False)
+        img2img_prompt, roll, img2img_prompt_style, img2img_negative_prompt, img2img_prompt_style2, submit, img2img_interrogate, img2img_prompt_style_apply, img2img_save_style, check_progress = create_toprow(is_img2img=True)
 
         with gr.Row().style(equal_height=False):
-
             with gr.Column(variant='panel'):
                 with gr.Group():
-                    switch_mode = gr.Radio(label='Mode', elem_id="img2img_mode", choices=['Redraw whole image', 'Inpaint a part of image', 'Loopback', 'SD upscale'], value='Redraw whole image', type="index", show_label=False)
+                    switch_mode = gr.Radio(label='Mode', elem_id="img2img_mode", choices=['Redraw whole image', 'Inpaint a part of image', 'SD upscale'], value='Redraw whole image', type="index", show_label=False)
                     init_img = gr.Image(label="Image for img2img", source="upload", interactive=True, type="pil")
-                    init_img_with_mask = gr.Image(label="Image for inpainting with mask", elem_id="img2maskimg", source="upload", interactive=True, type="pil", tool="sketch", visible=False)
-                    resize_mode = gr.Radio(label="Resize mode", show_label=False, choices=["Just resize", "Crop and resize", "Resize and fill"], type="index", value="Just resize")
+                    init_img_with_mask = gr.Image(label="Image for inpainting with mask", elem_id="img2maskimg", source="upload", interactive=True, type="pil", tool="sketch", visible=False, image_mode="RGBA")
+                    init_mask = gr.Image(label="Mask", source="upload", interactive=True, type="pil", visible=False)
+                    init_img_with_mask_comment = gr.HTML(elem_id="mask_bug_info", value="<small>if the editor shows ERROR, switch to another tab and back, then to another img2img mode above and back</small>", visible=False)
+
+                    with gr.Row():
+                        resize_mode = gr.Radio(label="Resize mode", elem_id="resize_mode", show_label=False, choices=["Just resize", "Crop and resize", "Resize and fill"], type="index", value="Just resize")
+                        mask_mode = gr.Radio(label="Mask mode", show_label=False, choices=["Draw mask", "Upload mask"], type="index", value="Draw mask")
 
                 steps = gr.Slider(minimum=1, maximum=150, step=1, label="Sampling Steps", value=20)
                 sampler_index = gr.Radio(label='Sampling method', choices=[x.name for x in samplers_for_img2img], value=samplers_for_img2img[0].name, type="index")
@@ -348,14 +448,14 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                     batch_size = gr.Slider(minimum=1, maximum=8, step=1, label='Batch size', value=1)
 
                 with gr.Group():
-                    cfg_scale = gr.Slider(minimum=1.0, maximum=15.0, step=0.5, label='CFG Scale', value=7.0)
-                    denoising_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Denoising Strength', value=0.75)
+                    cfg_scale = gr.Slider(minimum=1.0, maximum=30.0, step=0.5, label='CFG Scale', value=7.0)
+                    denoising_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Denoising strength', value=0.75)
 
                 with gr.Group():
-                    height = gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512)
                     width = gr.Slider(minimum=64, maximum=2048, step=64, label="Width", value=512)
+                    height = gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512)
 
-                seed = gr.Number(label='Seed', value=-1)
+                seed, subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w = create_seed_inputs()
 
                 with gr.Group():
                     custom_inputs = modules.scripts.scripts_img2img.setup_ui(is_img2img=True)
@@ -363,7 +463,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
             with gr.Column(variant='panel'):
                 with gr.Group():
                     img2img_preview = gr.Image(elem_id='img2img_preview', visible=False)
-                    img2img_gallery = gr.Gallery(label='Output', elem_id='img2img_gallery')
+                    img2img_gallery = gr.Gallery(label='Output', elem_id='img2img_gallery').style(grid=4)
 
                 with gr.Group():
                     with gr.Row():
@@ -372,6 +472,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                         img2img_send_to_inpaint = gr.Button('Send to inpaint')
                         img2img_send_to_extras = gr.Button('Send to extras')
                         interrupt = gr.Button('Interrupt')
+                        img2img_save_style = gr.Button('Save prompt as style')
 
                 progressbar = gr.HTML(elem_id="progressbar")
 
@@ -379,49 +480,71 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                     html_info = gr.HTML()
                     generation_info = gr.Textbox(visible=False)
 
-            def apply_mode(mode):
+            def apply_mode(mode, uploadmask):
                 is_classic = mode == 0
                 is_inpaint = mode == 1
-                is_loopback = mode == 2
-                is_upscale = mode == 3
+                is_upscale = mode == 2
 
                 return {
-                    init_img: gr_show(not is_inpaint),
-                    init_img_with_mask: gr_show(is_inpaint),
+                    init_img: gr_show(not is_inpaint or (is_inpaint and uploadmask == 1)),
+                    init_img_with_mask: gr_show(is_inpaint and uploadmask == 0),
+                    init_img_with_mask_comment: gr_show(is_inpaint and uploadmask == 0),
+                    init_mask: gr_show(is_inpaint and uploadmask == 1),
+                    mask_mode: gr_show(is_inpaint),
                     mask_blur: gr_show(is_inpaint),
                     inpainting_fill: gr_show(is_inpaint),
-                    batch_count: gr_show(not is_upscale),
-                    batch_size: gr_show(not is_loopback),
                     sd_upscale_upscaler_name: gr_show(is_upscale),
                     sd_upscale_overlap: gr_show(is_upscale),
                     inpaint_full_res: gr_show(is_inpaint),
                     inpainting_mask_invert: gr_show(is_inpaint),
+                    img2img_interrogate: gr_show(not is_inpaint),
                 }
 
             switch_mode.change(
                 apply_mode,
-                inputs=[switch_mode],
+                inputs=[switch_mode, mask_mode],
                 outputs=[
                     init_img,
                     init_img_with_mask,
+                    init_img_with_mask_comment,
+                    init_mask,
+                    mask_mode,
                     mask_blur,
                     inpainting_fill,
-                    batch_count,
-                    batch_size,
                     sd_upscale_upscaler_name,
                     sd_upscale_overlap,
                     inpaint_full_res,
                     inpainting_mask_invert,
+                    img2img_interrogate,
                 ]
             )
 
+            mask_mode.change(
+                lambda mode: {
+                    init_img: gr_show(mode == 1),
+                    init_img_with_mask: gr_show(mode == 0),
+                    init_mask: gr_show(mode == 1),
+                },
+                inputs=[mask_mode],
+                outputs=[
+                    init_img,
+                    init_img_with_mask,
+                    init_mask,
+                ],
+            )
+
             img2img_args = dict(
                 fn=img2img,
                 _js="submit",
                 inputs=[
-                    prompt,
+                    img2img_prompt,
+                    img2img_negative_prompt,
+                    img2img_prompt_style,
+                    img2img_prompt_style2,
                     init_img,
                     init_img_with_mask,
+                    init_mask,
+                    mask_mode,
                     steps,
                     sampler_index,
                     mask_blur,
@@ -434,6 +557,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                     cfg_scale,
                     denoising_strength,
                     seed,
+                    subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w,
                     height,
                     width,
                     resize_mode,
@@ -449,9 +573,15 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                 ]
             )
 
-            prompt.submit(**img2img_args)
+            img2img_prompt.submit(**img2img_args)
             submit.click(**img2img_args)
 
+            img2img_interrogate.click(
+                fn=interrogate,
+                inputs=[init_img],
+                outputs=[img2img_prompt],
+            )
+
             check_progress.click(
                 fn=check_progress_call,
                 show_progress=False,
@@ -467,9 +597,11 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
 
             save.click(
                 fn=wrap_gradio_call(save_files),
+                _js = "(x, y, z) => [x, y, selected_gallery_index()]",
                 inputs=[
                     generation_info,
                     img2img_gallery,
+                    html_info
                 ],
                 outputs=[
                     html_info,
@@ -478,33 +610,36 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                 ]
             )
 
-            send_to_img2img.click(
-                fn=lambda x: image_from_url_text(x),
-                _js="extract_image_from_gallery",
-                inputs=[txt2img_gallery],
-                outputs=[init_img],
+            roll.click(
+                fn=roll_artist,
+                inputs=[
+                    img2img_prompt,
+                ],
+                outputs=[
+                    img2img_prompt,
+                ]
             )
 
-            send_to_inpaint.click(
-                fn=lambda x: image_from_url_text(x),
-                _js="extract_image_from_gallery",
-                inputs=[txt2img_gallery],
-                outputs=[init_img_with_mask],
-            )
+            prompts = [(txt2img_prompt, txt2img_negative_prompt), (img2img_prompt, img2img_negative_prompt)]
+            style_dropdowns = [(txt2img_prompt_style, txt2img_prompt_style2), (img2img_prompt_style, img2img_prompt_style2)]
 
-            img2img_send_to_img2img.click(
-                fn=lambda x: image_from_url_text(x),
-                _js="extract_image_from_gallery",
-                inputs=[img2img_gallery],
-                outputs=[init_img],
-            )
+            dummy_component = gr.Label(visible=False)
+            for button, (prompt, negative_prompt) in zip([txt2img_save_style, img2img_save_style], prompts):
+                button.click(
+                    fn=add_style,
+                    _js="ask_for_style_name",
+                    # Have to pass empty dummy component here, because the JavaScript and Python function have to accept
+                    # the same number of parameters, but we only know the style-name after the JavaScript prompt
+                    inputs=[dummy_component, prompt, negative_prompt],
+                    outputs=[txt2img_prompt_style, img2img_prompt_style, txt2img_prompt_style2, img2img_prompt_style2],
+                )
 
-            img2img_send_to_inpaint.click(
-                fn=lambda x: image_from_url_text(x),
-                _js="extract_image_from_gallery",
-                inputs=[img2img_gallery],
-                outputs=[init_img_with_mask],
-            )
+            for button, (prompt, negative_prompt), (style1, style2) in zip([txt2img_prompt_style_apply, img2img_prompt_style_apply], prompts, style_dropdowns):
+                button.click(
+                    fn=apply_styles,
+                    inputs=[prompt, negative_prompt, style1, style2],
+                    outputs=[prompt, negative_prompt, style1, style2],
+                )
 
     with gr.Blocks(analytics_enabled=False) as extras_interface:
         with gr.Row().style(equal_height=False):
@@ -556,20 +691,6 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
 
         submit.click(**extras_args)
 
-        send_to_extras.click(
-            fn=lambda x: image_from_url_text(x),
-            _js="extract_image_from_gallery",
-            inputs=[txt2img_gallery],
-            outputs=[image],
-        )
-
-        img2img_send_to_extras.click(
-            fn=lambda x: image_from_url_text(x),
-            _js="extract_image_from_gallery",
-            inputs=[img2img_gallery],
-            outputs=[image],
-        )
-
     pnginfo_interface = gr.Interface(
         wrap_gradio_call(run_pnginfo),
         inputs=[
@@ -591,51 +712,66 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
         info = opts.data_labels[key]
         t = type(info.default)
 
+        args = info.component_args() if callable(info.component_args) else info.component_args
+
         if info.component is not None:
-            args = info.component_args() if callable(info.component_args) else info.component_args
-            item = info.component(label=info.label, value=fun, **(args or {}))
+            comp = info.component
         elif t == str:
-            item = gr.Textbox(label=info.label, value=fun, lines=1)
+            comp = gr.Textbox
         elif t == int:
-            item = gr.Number(label=info.label, value=fun)
+            comp = gr.Number
         elif t == bool:
-            item = gr.Checkbox(label=info.label, value=fun)
+            comp = gr.Checkbox
         else:
             raise Exception(f'bad options item type: {str(t)} for key {key}')
 
-        return item
+        return comp(label=info.label, value=fun, **(args or {}))
+
+    components = []
+    keys = list(opts.data_labels.keys())
+    settings_cols = 3
+    items_per_col = math.ceil(len(keys) / settings_cols)
 
     def run_settings(*args):
         up = []
 
-        for key, value, comp in zip(opts.data_labels.keys(), args, settings_interface.input_components):
+        for key, value, comp in zip(opts.data_labels.keys(), args, components):
+            comp_args = opts.data_labels[key].component_args
+            if comp_args and isinstance(comp_args, dict) and comp_args.get('visible') is False:
+                continue
+
             opts.data[key] = value
             up.append(comp.update(value=value))
 
         opts.save(shared.config_filename)
 
-        return 'Settings saved.', '', ''
+        return 'Settings applied.'
 
-    settings_interface = gr.Interface(
-        run_settings,
-        inputs=[create_setting_component(key) for key in opts.data_labels.keys()],
-        outputs=[
-            gr.Textbox(label='Result'),
-            gr.HTML(),
-            gr.HTML(),
-        ],
-        title=None,
-        description=None,
-        allow_flagging="never",
-        analytics_enabled=False,
-    )
+    with gr.Blocks(analytics_enabled=False) as settings_interface:
+        submit = gr.Button(value="Apply settings", variant='primary')
+        result = gr.HTML()
+
+        with gr.Row(elem_id="settings").style(equal_height=False):
+            for colno in range(settings_cols):
+                with gr.Column(variant='panel'):
+                    for rowno in range(items_per_col):
+                        index = rowno + colno * items_per_col
+
+                        if index < len(keys):
+                            components.append(create_setting_component(keys[index]))
+
+        submit.click(
+            fn=run_settings,
+            inputs=components,
+            outputs=[result]
+        )
 
     interfaces = [
-        (txt2img_interface, "txt2img"),
-        (img2img_interface, "img2img"),
-        (extras_interface, "Extras"),
-        (pnginfo_interface, "PNG Info"),
-        (settings_interface, "Settings"),
+        (txt2img_interface, "txt2img", "txt2img"),
+        (img2img_interface, "img2img", "img2img"),
+        (extras_interface, "Extras", "extras"),
+        (pnginfo_interface, "PNG Info", "pnginfo"),
+        (settings_interface, "Settings", "settings"),
     ]
 
     with open(os.path.join(script_path, "style.css"), "r", encoding="utf8") as file:
@@ -644,14 +780,61 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
     if not cmd_opts.no_progressbar_hiding:
         css += css_hide_progressbar
 
-    demo = gr.TabbedInterface(
-        interface_list=[x[0] for x in interfaces],
-        tab_names=[x[1] for x in interfaces],
-        analytics_enabled=False,
-        css=css,
-    )
+    with gr.Blocks(css=css, analytics_enabled=False, title="Stable Diffusion") as demo:
+        with gr.Tabs() as tabs:
+            for interface, label, ifid in interfaces:
+                with gr.TabItem(label, id=ifid):
+                    interface.render()
 
-    ui_config_file = os.path.join(modules.paths.script_path, 'ui-config.json')
+        tabs.change(
+            fn=lambda x: x,
+            inputs=[init_img_with_mask],
+            outputs=[init_img_with_mask],
+        )
+
+        send_to_img2img.click(
+            fn=lambda x: image_from_url_text(x),
+            _js="extract_image_from_gallery_img2img",
+            inputs=[txt2img_gallery],
+            outputs=[init_img],
+        )
+
+        send_to_inpaint.click(
+            fn=lambda x: image_from_url_text(x),
+            _js="extract_image_from_gallery_img2img",
+            inputs=[txt2img_gallery],
+            outputs=[init_img_with_mask],
+        )
+
+        img2img_send_to_img2img.click(
+            fn=lambda x: image_from_url_text(x),
+            _js="extract_image_from_gallery",
+            inputs=[img2img_gallery],
+            outputs=[init_img],
+        )
+
+        img2img_send_to_inpaint.click(
+            fn=lambda x: image_from_url_text(x),
+            _js="extract_image_from_gallery",
+            inputs=[img2img_gallery],
+            outputs=[init_img_with_mask],
+        )
+
+        send_to_extras.click(
+            fn=lambda x: image_from_url_text(x),
+            _js="extract_image_from_gallery_extras",
+            inputs=[txt2img_gallery],
+            outputs=[image],
+        )
+
+        img2img_send_to_extras.click(
+            fn=lambda x: image_from_url_text(x),
+            _js="extract_image_from_gallery_extras",
+            inputs=[img2img_gallery],
+            outputs=[image],
+        )
+
+    ui_config_file = cmd_opts.ui_config_file
     ui_settings = {}
     settings_count = len(ui_settings)
     error_loading = False
@@ -686,6 +869,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
 
     visit(txt2img_interface, loadsave, "txt2img")
     visit(img2img_interface, loadsave, "img2img")
+    visit(extras_interface, loadsave, "extras")
 
     if not error_loading and (not os.path.exists(ui_config_file) or settings_count != len(ui_settings)):
         with open(ui_config_file, "w", encoding="utf8") as file:

requirements.txt

@@ -1,3 +1,5 @@
+transformers
+diffusers
 basicsr
 gfpgan
 gradio
@@ -10,5 +12,11 @@ omegaconf
 pytorch_lightning
 diffusers
 invisible-watermark
+scikit-image>=0.19
+fonts
+font-roboto
 git+https://github.com/crowsonkb/k-diffusion.git
-git+https://github.com/TencentARC/GFPGAN.git
+git+https://github.com/TencentARC/GFPGAN.git@8d2447a2d918f8eba5a4a01463fd48e45126a379
+timm==0.4.12
+fairscale==0.4.4
+piexif

requirements_versions.txt

@@ -1,10 +1,18 @@
+transformers==4.19.2
+diffusers==0.2.4
 basicsr==1.3.5
 gfpgan
-gradio==3.2
-numpy==1.22.0
+gradio==3.3
+numpy==1.23.3
 Pillow==9.2.0
 realesrgan==0.2.5.0
 torch
 transformers==4.19.2
 omegaconf==2.1.1
 pytorch_lightning==1.7.2
+scikit-image==0.19.2
+fonts
+font-roboto
+timm==0.4.12
+fairscale==0.4.4
+piexif==1.1.3

script.js

@@ -1,8 +1,10 @@
+
+
 titles = {
-    "Sampling steps": "How many times to imptove the generated image itratively; higher values take longer; very low values can produce bad results",
+    "Sampling steps": "How many times to improve the generated image iteratively; higher values take longer; very low values can produce bad results",
     "Sampling method": "Which algorithm to use to produce the image",
 	"GFPGAN": "Restore low quality faces using GFPGAN neural network",
-	"Euler a": "Euler Ancestral - very creative, each can get acompletely different pictures depending on step count, setting seps tohigher than 30-40 does not help",
+	"Euler a": "Euler Ancestral - very creative, each can get a completely different picture depending on step count, setting steps to higher than 30-40 does not help",
 	"DDIM": "Denoising Diffusion Implicit Models - best at inpainting",
 
 	"Batch count": "How many batches of images to create",
@@ -11,7 +13,6 @@ titles = {
     "Seed": "A value that determines the output of random number generator - if you create an image with same parameters and seed as another image, you'll get the same result",
 
     "Inpaint a part of image": "Draw a mask over an image, and the script will regenerate the masked area with content according to prompt",
-    "Loopback": "Process an image, use it as an input, repeat. Batch count determings number of iterations.",
     "SD upscale": "Upscale image normally, split result into tiles, improve each tile using img2img, merge whole image back",
 
     "Just resize": "Resize image to target resolution. Unless height and width match, you will get incorrect aspect ratio.",
@@ -26,7 +27,8 @@ titles = {
     "latent nothing": "fill it with latent space zeroes",
     "Inpaint at full resolution": "Upscale masked region to target resolution, do inpainting, downscale back and paste into original image",
 
-    "Denoising Strength": "Determines how little respect the algorithm should have for image's content. At 0, nothing will change, and at 1 you'll get an unrelated image.",
+    "Denoising strength": "Determines how little respect the algorithm should have for image's content. At 0, nothing will change, and at 1 you'll get an unrelated image. With values below 1.0, processing will take less steps than the Sampling Steps slider specifies.",
+    "Denoising strength change factor": "In loopback mode, on each loop the denoising strength is multiplied by this value. <1 means decreasing variety so your sequence will converge on a fixed picture. >1 means increasing variety so your sequence will become more and more chaotic.",
 
     "Interrupt": "Stop processing images and return any results accumulated so far.",
     "Save": "Write image to a directory (default - log/images) and generation parameters into csv file.",
@@ -36,15 +38,34 @@ titles = {
 
     "None": "Do not do anything special",
     "Prompt matrix": "Separate prompts into parts using vertical pipe character (|) and the script will create a picture for every combination of them (except for the first part, which will be present in all combinations)",
-    "X/Y plot": "Create a grid where images will have different parameters. Use inputs below to specify which parameterswill be shared by columns and rows",
-    "Custom code": "Run python code. Advanced user only. Must run program with --allow-code for this to work",
+    "X/Y plot": "Create a grid where images will have different parameters. Use inputs below to specify which parameters will be shared by columns and rows",
+    "Custom code": "Run Python code. Advanced user only. Must run program with --allow-code for this to work",
 
     "Prompt S/R": "Separate a list of words with commas, and the first word will be used as a keyword: script will search for this word in the prompt, and replace it with others",
 
     "Tiling": "Produce an image that can be tiled.",
-    "Tile overlap": "For SD upscale, how much overlap in pixels should there be between tiles. Tils overlap so that when they are merged back into one oicture, there is no clearly visible seam.",
+    "Tile overlap": "For SD upscale, how much overlap in pixels should there be between tiles. Tiles overlap so that when they are merged back into one picture, there is no clearly visible seam.",
 
     "Roll": "Add a random artist to the prompt.",
+
+    "Variation seed": "Seed of a different picture to be mixed into the generation.",
+    "Variation strength": "How strong of a variation to produce. At 0, there will be no effect. At 1, you will get the complete picture with variation seed (except for ancestral samplers, where you will just get something).",
+    "Resize seed from height": "Make an attempt to produce a picture similar to what would have been produced with same seed at specified resolution",
+    "Resize seed from width": "Make an attempt to produce a picture similar to what would have been produced with same seed at specified resolution",
+
+    "Interrogate": "Reconstruct prompt from existing image and put it into the prompt field.",
+
+    "Images filename pattern": "Use following tags to define how filenames for images are chosen: [steps], [cfg], [prompt], [prompt_spaces], [width], [height], [sampler], [seed], [model_hash], [prompt_words], [date]; leave empty for default.",
+    "Directory name pattern": "Use following tags to define how subdirectories for images and grids are chosen: [steps], [cfg], [prompt], [prompt_spaces], [width], [height], [sampler], [seed], [model_hash], [prompt_words], [date]; leave empty for default.",
+
+    "Loopback": "Process an image, use it as an input, repeat.",
+    "Loops": "How many times to repeat processing an image and using it as input for the next iteration",
+
+
+    "Style 1": "Style to apply; styles have components for both positive and negative prompts and apply to both",
+    "Style 2": "Style to apply; styles have components for both positive and negative prompts and apply to both",
+    "Apply style": "Insert selected styles into prompt fields",
+    "Create style": "Save current prompts as a style. If you add the token {prompt} to the text, the style use that as placeholder for your prompt when you use the style in the future.",
 }
 
 function gradioApp(){
@@ -129,12 +150,21 @@ function extract_image_from_gallery(gallery){
     index = selected_gallery_index()
 
     if (index < 0 || index >= gallery.length){
-        return []
+        return [null]
     }
 
     return gallery[index];
 }
 
+function extract_image_from_gallery_img2img(gallery){
+    gradioApp().querySelectorAll('button')[1].click();
+    return extract_image_from_gallery(gallery);
+}
+
+function extract_image_from_gallery_extras(gallery){
+    gradioApp().querySelectorAll('button')[2].click();
+    return extract_image_from_gallery(gallery);
+}
 
 function requestProgress(){
     btn = gradioApp().getElementById("check_progress");
@@ -150,6 +180,26 @@ function submit(){
     for(var i=0;i<arguments.length;i++){
         res.push(arguments[i])
     }
-    console.log(res)
     return res
-}
+}
+
+window.addEventListener('paste', e => {
+    const files = e.clipboardData.files;
+    if (!files || files.length !== 1) {
+        return;
+    }
+    if (!['image/png', 'image/gif', 'image/jpeg'].includes(files[0].type)) {
+        return;
+    }
+    [...gradioApp().querySelectorAll('input[type=file][accept="image/x-png,image/gif,image/jpeg"]')]
+        .filter(input => !input.matches('.\\!hidden input[type=file]'))
+        .forEach(input => {
+            input.files = files;
+            input.dispatchEvent(new Event('change'))
+        });
+});
+
+function ask_for_style_name(_, prompt_text, negative_prompt_text) {
+    name_ = prompt('Style name:')
+    return name_ === null ? [null, null, null]: [name_, prompt_text, negative_prompt_text]
+}

scripts/batch.py

@@ -0,0 +1,59 @@
+import math
+import os
+import sys
+import traceback
+
+import modules.scripts as scripts
+import gradio as gr
+
+from modules.processing import Processed, process_images
+from PIL import Image
+from modules.shared import opts, cmd_opts, state
+
+
+class Script(scripts.Script):
+    def title(self):
+        return "Batch processing"
+
+    def show(self, is_img2img):
+        return is_img2img
+
+    def ui(self, is_img2img):
+        input_dir = gr.Textbox(label="Input directory", lines=1)
+        output_dir = gr.Textbox(label="Output directory", lines=1)
+
+        return [input_dir, output_dir]
+
+    def run(self, p, input_dir, output_dir):
+        images = [file for file in [os.path.join(input_dir, x) for x in os.listdir(input_dir)] if os.path.isfile(file)]
+
+        batch_count = math.ceil(len(images) / p.batch_size)
+        print(f"Will process {len(images)} images in {batch_count} batches.")
+
+        p.batch_count = 1
+        p.do_not_save_grid = True
+        p.do_not_save_samples = True
+
+        state.job_count = batch_count
+
+        for batch_no in range(batch_count):
+            batch_images = []
+            for path in images[batch_no*p.batch_size:(batch_no+1)*p.batch_size]:
+                try:
+                    img = Image.open(path)
+                    batch_images.append((img, path))
+                except:
+                    print(f"Error processing {path}:", file=sys.stderr)
+                    print(traceback.format_exc(), file=sys.stderr)
+
+            if len(batch_images) == 0:
+                continue
+
+            state.job = f"{batch_no} out of {batch_count}: {batch_images[0][1]}"
+            p.init_images = [x[0] for x in batch_images]
+            proc = process_images(p)
+            for image, (_, path) in zip(proc.images, batch_images):
+                filename = os.path.basename(path)
+                image.save(os.path.join(output_dir, filename))
+
+        return Processed(p, [], p.seed, "")

scripts/img2imgalt.py

@@ -0,0 +1,116 @@
+from collections import namedtuple
+
+import numpy as np
+from tqdm import trange
+
+import modules.scripts as scripts
+import gradio as gr
+
+from modules import processing, shared, sd_samplers
+from modules.processing import Processed
+from modules.sd_samplers import samplers
+from modules.shared import opts, cmd_opts, state
+
+import torch
+import k_diffusion as K
+
+from PIL import Image
+from torch import autocast
+from einops import rearrange, repeat
+
+
+def find_noise_for_image(p, cond, uncond, cfg_scale, steps):
+    x = p.init_latent
+
+    s_in = x.new_ones([x.shape[0]])
+    dnw = K.external.CompVisDenoiser(shared.sd_model)
+    sigmas = dnw.get_sigmas(steps).flip(0)
+
+    shared.state.sampling_steps = steps
+
+    for i in trange(1, len(sigmas)):
+        shared.state.sampling_step += 1
+
+        x_in = torch.cat([x] * 2)
+        sigma_in = torch.cat([sigmas[i] * s_in] * 2)
+        cond_in = torch.cat([uncond, cond])
+
+        c_out, c_in = [K.utils.append_dims(k, x_in.ndim) for k in dnw.get_scalings(sigma_in)]
+        t = dnw.sigma_to_t(sigma_in)
+
+        eps = shared.sd_model.apply_model(x_in * c_in, t, cond=cond_in)
+        denoised_uncond, denoised_cond = (x_in + eps * c_out).chunk(2)
+
+        denoised = denoised_uncond + (denoised_cond - denoised_uncond) * cfg_scale
+
+        d = (x - denoised) / sigmas[i]
+        dt = sigmas[i] - sigmas[i - 1]
+
+        x = x + d * dt
+
+        sd_samplers.store_latent(x)
+
+        # This shouldn't be necessary, but solved some VRAM issues
+        del x_in, sigma_in, cond_in, c_out, c_in, t,
+        del eps, denoised_uncond, denoised_cond, denoised, d, dt
+
+    shared.state.nextjob()
+
+    return x / x.std()
+
+
+Cached = namedtuple("Cached", ["noise", "cfg_scale", "steps", "latent", "original_prompt"])
+
+
+class Script(scripts.Script):
+    def __init__(self):
+        self.cache = None
+
+    def title(self):
+        return "img2img alternative test"
+
+    def show(self, is_img2img):
+        return is_img2img
+
+    def ui(self, is_img2img):
+        original_prompt = gr.Textbox(label="Original prompt", lines=1)
+        cfg = gr.Slider(label="Decode CFG scale", minimum=0.0, maximum=15.0, step=0.1, value=1.0)
+        st = gr.Slider(label="Decode steps", minimum=1, maximum=150, step=1, value=50)
+
+        return [original_prompt, cfg, st]
+
+    def run(self, p, original_prompt, cfg, st):
+        p.batch_size = 1
+        p.batch_count = 1
+
+        def sample_extra(x, conditioning, unconditional_conditioning):
+            lat = (p.init_latent.cpu().numpy() * 10).astype(int)
+
+            same_params = self.cache is not None and self.cache.cfg_scale == cfg and self.cache.steps == st and self.cache.original_prompt == original_prompt
+            same_everything = same_params and self.cache.latent.shape == lat.shape and np.abs(self.cache.latent-lat).sum() < 100
+
+            if same_everything:
+                noise = self.cache.noise
+            else:
+                shared.state.job_count += 1
+                cond = p.sd_model.get_learned_conditioning(p.batch_size * [original_prompt])
+                noise = find_noise_for_image(p, cond, unconditional_conditioning, cfg, st)
+                self.cache = Cached(noise, cfg, st, lat, original_prompt)
+
+            sampler = samplers[p.sampler_index].constructor(p.sd_model)
+
+            samples_ddim = sampler.sample(p, noise, conditioning, unconditional_conditioning)
+            return samples_ddim
+
+        p.sample = sample_extra
+
+        p.extra_generation_params = {
+            "Decode prompt": original_prompt,
+            "Decode CFG scale": cfg,
+            "Decode steps": st,
+        }
+
+        processed = processing.process_images(p)
+
+        return processed
+

scripts/loopback.py

@@ -0,0 +1,78 @@
+import numpy as np
+from tqdm import trange
+
+import modules.scripts as scripts
+import gradio as gr
+
+from modules import processing, shared, sd_samplers, images
+from modules.processing import Processed
+from modules.sd_samplers import samplers
+from modules.shared import opts, cmd_opts, state
+
+class Script(scripts.Script):
+    def title(self):
+        return "Loopback"
+
+    def show(self, is_img2img):
+        return is_img2img
+
+    def ui(self, is_img2img):
+        loops = gr.Slider(minimum=1, maximum=32, step=1, label='Loops', value=4)
+        denoising_strength_change_factor = gr.Slider(minimum=0.9, maximum=1.1, step=0.01, label='Denoising strength change factor', value=1)
+
+        return [loops, denoising_strength_change_factor]
+
+    def run(self, p, loops, denoising_strength_change_factor):
+        processing.fix_seed(p)
+        batch_count = p.n_iter
+        p.extra_generation_params = {
+            "Denoising strength change factor": denoising_strength_change_factor,
+        }
+
+        p.batch_size = 1
+        p.n_iter = 1
+
+        output_images, info = None, None
+        initial_seed = None
+        initial_info = None
+
+        grids = []
+        all_images = []
+        state.job_count = loops * batch_count
+
+        for n in range(batch_count):
+            history = []
+
+            for i in range(loops):
+                p.n_iter = 1
+                p.batch_size = 1
+                p.do_not_save_grid = True
+
+                state.job = f"Iteration {i + 1}/{loops}, batch {n + 1}/{batch_count}"
+
+                processed = processing.process_images(p)
+
+                if initial_seed is None:
+                    initial_seed = processed.seed
+                    initial_info = processed.info
+
+                init_img = processed.images[0]
+
+                p.init_images = [init_img]
+                p.seed = processed.seed + 1
+                p.denoising_strength = min(max(p.denoising_strength * denoising_strength_change_factor, 0.1), 1)
+                history.append(processed.images[0])
+
+            grid = images.image_grid(history, rows=1)
+            if opts.grid_save:
+                images.save_image(grid, p.outpath_grids, "grid", initial_seed, p.prompt, opts.grid_format, info=info, short_filename=not opts.grid_extended_filename, grid=True, p=p)
+
+            grids.append(grid)
+            all_images += history
+
+        if opts.return_grid:
+            all_images = grids + all_images
+
+        processed = Processed(p, all_images, initial_seed, initial_info)
+
+        return processed

scripts/poor_mans_outpainting.py

@@ -4,7 +4,7 @@ import modules.scripts as scripts
 import gradio as gr
 from PIL import Image, ImageDraw
 
-from modules import images, processing
+from modules import images, processing, devices
 from modules.processing import Processed, process_images
 from modules.shared import opts, cmd_opts, state
 
@@ -21,7 +21,7 @@ class Script(scripts.Script):
         if not is_img2img:
             return None
 
-        pixels = gr.Slider(label="Pixels to expand", minimum=8, maximum=128, step=8)
+        pixels = gr.Slider(label="Pixels to expand", minimum=8, maximum=256, step=8, value=128)
         mask_blur = gr.Slider(label='Mask blur', minimum=0, maximum=64, step=1, value=4, visible=False)
         inpainting_fill = gr.Radio(label='Masked content', choices=['fill', 'original', 'latent noise', 'latent nothing'], value='fill', type="index", visible=False)
         direction = gr.CheckboxGroup(label="Outpainting direction", choices=['left', 'right', 'up', 'down'], value=['left', 'right', 'up', 'down'])
@@ -32,7 +32,7 @@ class Script(scripts.Script):
         initial_seed = None
         initial_info = None
 
-        p.mask_blur = mask_blur
+        p.mask_blur = mask_blur * 2
         p.inpainting_fill = inpainting_fill
         p.inpaint_full_res = False
 
@@ -47,11 +47,14 @@ class Script(scripts.Script):
 
         if left > 0:
             left = left * (target_w - init_img.width) // (left + right)
-        right = target_w - init_img.width - left
+        if right > 0:
+            right = target_w - init_img.width - left
 
         if up > 0:
             up = up * (target_h - init_img.height) // (up + down)
-        down = target_h - init_img.height - up
+
+        if down > 0:
+            down = target_h - init_img.height - up
 
         img = Image.new("RGB", (target_w, target_h))
         img.paste(init_img, (left, up))
@@ -67,13 +70,18 @@ class Script(scripts.Script):
 
         latent_mask = Image.new("L", (img.width, img.height), "white")
         latent_draw = ImageDraw.Draw(latent_mask)
-        latent_draw.rectangle((left + left//2, up + up//2, mask.width - right - right//2, mask.height - down - down//2), fill="black")
+        latent_draw.rectangle((
+             left + (mask_blur//2 if left > 0 else 0),
+             up + (mask_blur//2 if up > 0 else 0),
+             mask.width - right - (mask_blur//2 if right > 0 else 0),
+             mask.height - down - (mask_blur//2 if down > 0 else 0)
+        ), fill="black")
 
-        processing.torch_gc()
+        devices.torch_gc()
 
         grid = images.split_grid(img, tile_w=p.width, tile_h=p.height, overlap=pixels)
         grid_mask = images.split_grid(mask, tile_w=p.width, tile_h=p.height, overlap=pixels)
-        grid_latent_mask = images.split_grid(mask, tile_w=p.width, tile_h=p.height, overlap=pixels)
+        grid_latent_mask = images.split_grid(latent_mask, tile_w=p.width, tile_h=p.height, overlap=pixels)
 
         p.n_iter = 1
         p.batch_size = 1
@@ -85,8 +93,13 @@ class Script(scripts.Script):
         work_latent_mask = []
         work_results = []
 
-        for (_, _, row), (_, _, row_mask), (_, _, row_latent_mask) in zip(grid.tiles, grid_mask.tiles, grid_latent_mask.tiles):
+        for (y, h, row), (_, _, row_mask), (_, _, row_latent_mask) in zip(grid.tiles, grid_mask.tiles, grid_latent_mask.tiles):
             for tiledata, tiledata_mask, tiledata_latent_mask in zip(row, row_mask, row_latent_mask):
+                x, w = tiledata[0:2]
+
+                if x >= left and x+w <= img.width - right and y >= up and y+h <= img.height - down:
+                    continue
+
                 work.append(tiledata[2])
                 work_mask.append(tiledata_mask[2])
                 work_latent_mask.append(tiledata_latent_mask[2])
@@ -115,13 +128,18 @@ class Script(scripts.Script):
         image_index = 0
         for y, h, row in grid.tiles:
             for tiledata in row:
+                x, w = tiledata[0:2]
+
+                if x >= left and x+w <= img.width - right and y >= up and y+h <= img.height - down:
+                    continue
+
                 tiledata[2] = work_results[image_index] if image_index < len(work_results) else Image.new("RGB", (p.width, p.height))
                 image_index += 1
 
         combined_image = images.combine_grid(grid)
 
         if opts.samples_save:
-            images.save_image(combined_image, p.outpath_samples, "", initial_seed, p.prompt, opts.grid_format, info=initial_info)
+            images.save_image(combined_image, p.outpath_samples, "", initial_seed, p.prompt, opts.grid_format, info=initial_info, p=p)
 
         processed = Processed(p, [combined_image], initial_seed, initial_info)
 

scripts/prompt_matrix.py

@@ -50,7 +50,7 @@ class Script(scripts.Script):
         return [put_at_start]
 
     def run(self, p, put_at_start):
-        seed = modules.processing.set_seed(p.seed)
+        modules.processing.fix_seed(p)
 
         original_prompt = p.prompt[0] if type(p.prompt) == list else p.prompt
 
@@ -73,8 +73,8 @@ class Script(scripts.Script):
         print(f"Prompt matrix will create {len(all_prompts)} images using a total of {p.n_iter} batches.")
 
         p.prompt = all_prompts
+        p.seed = [p.seed for _ in all_prompts]
         p.prompt_for_display = original_prompt
-        p.seed = len(all_prompts) * [seed]
         processed = process_images(p)
 
         grid = images.image_grid(processed.images, p.batch_size, rows=1 << ((len(prompt_matrix_parts) - 1) // 2))
@@ -82,6 +82,6 @@ class Script(scripts.Script):
         processed.images.insert(0, grid)
 
         if opts.grid_save:
-            images.save_image(processed.images[0], p.outpath_grids, "prompt_matrix", prompt=original_prompt, seed=seed)
+            images.save_image(processed.images[0], p.outpath_grids, "prompt_matrix", prompt=original_prompt, seed=processed.seed, grid=True, p=p)
 
         return processed

scripts/prompts_from_file.py

@@ -0,0 +1,41 @@
+import math
+import os
+import sys
+import traceback
+
+import modules.scripts as scripts
+import gradio as gr
+
+from modules.processing import Processed, process_images
+from PIL import Image
+from modules.shared import opts, cmd_opts, state
+
+
+class Script(scripts.Script):
+    def title(self):
+        return "Prompts from file"
+
+    def ui(self, is_img2img):
+        file = gr.File(label="File with inputs", type='bytes')
+
+        return [file]
+
+    def run(self, p, data: bytes):
+        lines = [x.strip() for x in data.decode('utf8', errors='ignore').split("\n")]
+        lines = [x for x in lines if len(x) > 0]
+
+        batch_count = math.ceil(len(lines) / p.batch_size)
+        print(f"Will process {len(lines) * p.n_iter} images in {batch_count * p.n_iter} batches.")
+
+        p.do_not_save_grid = True
+
+        state.job_count = batch_count
+
+        images = []
+        for batch_no in range(batch_count):
+            state.job = f"{batch_no} out of {batch_count * p.n_iter}"
+            p.prompt = lines[batch_no*p.batch_size:(batch_no+1)*p.batch_size] * p.n_iter
+            proc = process_images(p)
+            images += proc.images
+
+        return Processed(p, images, p.seed, "")

scripts/xy_grid.py

@@ -2,6 +2,8 @@ from collections import namedtuple
 from copy import copy
 import random
 
+import numpy as np
+
 import modules.scripts as scripts
 import gradio as gr
 
@@ -21,6 +23,7 @@ def apply_field(field):
 
 def apply_prompt(p, x, xs):
     p.prompt = p.prompt.replace(xs[0], x)
+    p.negative_prompt = p.negative_prompt.replace(xs[0], x)
 
 
 samplers_dict = {}
@@ -39,28 +42,43 @@ def apply_sampler(p, x, xs):
 
 
 def format_value_add_label(p, opt, x):
+    if type(x) == float:
+        x = round(x, 8)
+
     return f"{opt.label}: {x}"
 
 
 def format_value(p, opt, x):
+    if type(x) == float:
+        x = round(x, 8)
+
     return x
 
+def do_nothing(p, x, xs):
+    pass
+
+def format_nothing(p, opt, x):
+    return ""
+
 
 AxisOption = namedtuple("AxisOption", ["label", "type", "apply", "format_value"])
 AxisOptionImg2Img = namedtuple("AxisOptionImg2Img", ["label", "type", "apply", "format_value"])
 
 
 axis_options = [
+    AxisOption("Nothing", str, do_nothing, format_nothing),
     AxisOption("Seed", int, apply_field("seed"), format_value_add_label),
+    AxisOption("Var. seed", int, apply_field("subseed"), format_value_add_label),
+    AxisOption("Var. strength", float, apply_field("subseed_strength"), format_value_add_label),
     AxisOption("Steps", int, apply_field("steps"), format_value_add_label),
     AxisOption("CFG Scale", float, apply_field("cfg_scale"), format_value_add_label),
     AxisOption("Prompt S/R", str, apply_prompt, format_value),
     AxisOption("Sampler", str, apply_sampler, format_value),
-    AxisOptionImg2Img("Denoising", float, apply_field("denoising_strength"), format_value_add_label) #  as it is now all AxisOptionImg2Img items must go after AxisOption ones
+    AxisOptionImg2Img("Denoising", float, apply_field("denoising_strength"), format_value_add_label), #  as it is now all AxisOptionImg2Img items must go after AxisOption ones
 ]
 
 
-def draw_xy_grid(xs, ys, x_label, y_label, cell):
+def draw_xy_grid(p, xs, ys, x_label, y_label, cell, draw_legend):
     res = []
 
     ver_texts = [[images.GridAnnotation(y_label(y))] for y in ys]
@@ -68,7 +86,7 @@ def draw_xy_grid(xs, ys, x_label, y_label, cell):
 
     first_pocessed = None
 
-    state.job_count = len(xs) * len(ys)
+    state.job_count = len(xs) * len(ys) * p.n_iter
 
     for iy, y in enumerate(ys):
         for ix, x in enumerate(xs):
@@ -81,7 +99,8 @@ def draw_xy_grid(xs, ys, x_label, y_label, cell):
             res.append(processed.images[0])
 
     grid = images.image_grid(res, rows=len(ys))
-    grid = images.draw_grid_annotations(grid, res[0].width, res[0].height, hor_texts, ver_texts)
+    if draw_legend:
+        grid = images.draw_grid_annotations(grid, res[0].width, res[0].height, hor_texts, ver_texts)
 
     first_pocessed.images = [grid]
 
@@ -89,6 +108,10 @@ def draw_xy_grid(xs, ys, x_label, y_label, cell):
 
 
 re_range = re.compile(r"\s*([+-]?\s*\d+)\s*-\s*([+-]?\s*\d+)(?:\s*\(([+-]\d+)\s*\))?\s*")
+re_range_float = re.compile(r"\s*([+-]?\s*\d+(?:.\d*)?)\s*-\s*([+-]?\s*\d+(?:.\d*)?)(?:\s*\(([+-]\d+(?:.\d*)?)\s*\))?\s*")
+
+re_range_count = re.compile(r"\s*([+-]?\s*\d+)\s*-\s*([+-]?\s*\d+)(?:\s*\[(\d+)\s*\])?\s*")
+re_range_count_float = re.compile(r"\s*([+-]?\s*\d+(?:.\d*)?)\s*-\s*([+-]?\s*\d+(?:.\d*)?)(?:\s*\[(\d+(?:.\d*)?)\s*\])?\s*")
 
 class Script(scripts.Script):
     def title(self):
@@ -98,19 +121,20 @@ class Script(scripts.Script):
         current_axis_options = [x for x in axis_options if type(x) == AxisOption or type(x) == AxisOptionImg2Img and is_img2img]
 
         with gr.Row():
-            x_type = gr.Dropdown(label="X type", choices=[x.label for x in current_axis_options], value=current_axis_options[0].label, visible=False, type="index", elem_id="x_type")
+            x_type = gr.Dropdown(label="X type", choices=[x.label for x in current_axis_options], value=current_axis_options[1].label, visible=False, type="index", elem_id="x_type")
             x_values = gr.Textbox(label="X values", visible=False, lines=1)
 
         with gr.Row():
-            y_type = gr.Dropdown(label="Y type", choices=[x.label for x in current_axis_options], value=current_axis_options[1].label, visible=False, type="index", elem_id="y_type")
+            y_type = gr.Dropdown(label="Y type", choices=[x.label for x in current_axis_options], value=current_axis_options[4].label, visible=False, type="index", elem_id="y_type")
             y_values = gr.Textbox(label="Y values", visible=False, lines=1)
+        
+        draw_legend = gr.Checkbox(label='Draw legend', value=True)
+            
+        return [x_type, x_values, y_type, y_values, draw_legend]
 
-        return [x_type, x_values, y_type, y_values]
-
-    def run(self, p, x_type, x_values, y_type, y_values):
-        p.seed = modules.processing.set_seed(p.seed)
+    def run(self, p, x_type, x_values, y_type, y_values, draw_legend):
+        modules.processing.fix_seed(p)
         p.batch_size = 1
-        p.batch_count = 1
 
         def process_axis(opt, vals):
             valslist = [x.strip() for x in vals.split(",")]
@@ -120,6 +144,7 @@ class Script(scripts.Script):
 
                 for val in valslist:
                     m = re_range.fullmatch(val)
+                    mc = re_range_count.fullmatch(val)
                     if m is not None:
 
                         start = int(m.group(1))
@@ -127,6 +152,34 @@ class Script(scripts.Script):
                         step = int(m.group(3)) if m.group(3) is not None else 1
 
                         valslist_ext += list(range(start, end, step))
+                    elif mc is not None:
+                        start = int(mc.group(1))
+                        end   = int(mc.group(2))
+                        num   = int(mc.group(3)) if mc.group(3) is not None else 1
+                        
+                        valslist_ext += [int(x) for x in np.linspace(start = start, stop = end, num = num).tolist()]
+                    else:
+                        valslist_ext.append(val)
+
+                valslist = valslist_ext
+            elif opt.type == float:
+                valslist_ext = []
+
+                for val in valslist:
+                    m = re_range_float.fullmatch(val)
+                    mc = re_range_count_float.fullmatch(val)
+                    if m is not None:
+                        start = float(m.group(1))
+                        end = float(m.group(2))
+                        step = float(m.group(3)) if m.group(3) is not None else 1
+
+                        valslist_ext += np.arange(start, end + step, step).tolist()
+                    elif mc is not None:
+                        start = float(mc.group(1))
+                        end   = float(mc.group(2))
+                        num   = int(mc.group(3)) if mc.group(3) is not None else 1
+                        
+                        valslist_ext += np.linspace(start = start, stop = end, num = num).tolist()
                     else:
                         valslist_ext.append(val)
 
@@ -150,14 +203,16 @@ class Script(scripts.Script):
             return process_images(pc)
 
         processed = draw_xy_grid(
+            p,
             xs=xs,
             ys=ys,
             x_label=lambda x: x_opt.format_value(p, x_opt, x),
             y_label=lambda y: y_opt.format_value(p, y_opt, y),
-            cell=cell
+            cell=cell,
+            draw_legend=draw_legend
         )
 
         if opts.grid_save:
-            images.save_image(processed.images[0], p.outpath_grids, "xy_grid", prompt=p.prompt, seed=processed.seed)
+            images.save_image(processed.images[0], p.outpath_grids, "xy_grid", prompt=p.prompt, seed=processed.seed, grid=True, p=p)
 
         return processed

style.css

@@ -1,23 +1,70 @@
 .output-html p {margin: 0 0.5em;}
 .performance { font-size: 0.85em; color: #444; }
 
-#txt2img_generate, #img2img_generate{
-    max-width: 13em;
+#generate{
+    min-height: 4.5em;
 }
 
-#txt2img_roll{
+#txt2img_gallery, #img2img_gallery{
+    min-height: 768px;
+}
+#txt2img_gallery img, #img2img_gallery img{
+    object-fit: scale-down;
+}
+
+#subseed_show{
+    min-width: 6em;
+    max-width: 6em;
+}
+
+#subseed_show label{
+    height: 100%;
+}
+
+#roll{
     min-width: 1em;
     max-width: 4em;
+    margin: 0.5em;
+}
+
+#style_apply, #style_create, #interrogate{
+    margin: 0.75em 0.25em 0.25em 0.25em;
+    min-width: 3em;
+}
+
+#style_pos_col, #style_neg_col{
+    min-width: 4em !important;
+}
+
+#style_index, #style2_index{
+    margin-top: 1em;
+}
+
+.gr-form{
+    background: transparent;
+}
+
+#toprow div{
+    border: none;
+    gap: 0;
+    background: transparent;
+}
+
+#resize_mode{
+    flex: 1.5;
 }
 
 button{
     align-self: stretch !important;
 }
 
-#img2img_prompt, #txt2img_prompt{
-    padding: 0;
+#prompt, #negative_prompt{
     border: none !important;
 }
+#prompt textarea, #negative_prompt textarea{
+    border: none !important;
+}
+
 
 #img2maskimg .h-60{
     height: 30rem;
@@ -65,6 +112,10 @@ fieldset span.text-gray-500, .gr-block.gr-box span.text-gray-500,  label.block s
     border-right: 1px solid rgb(55 65 81);
 }
 
+#settings fieldset span.text-gray-500, #settings .gr-block.gr-box span.text-gray-500, #settings label.block span{
+    position: relative;
+    border: none;
+}
 
 .gr-panel div.flex-col div.justify-between label span{
     margin: 0;
@@ -77,6 +128,11 @@ fieldset span.text-gray-500, .gr-block.gr-box span.text-gray-500,  label.block s
   padding: 0 0.5em;
 }
 
+#settings .gr-panel div.flex-col div.justify-between div{
+    position: relative;
+    z-index: 200;
+}
+
 input[type="range"]{
     margin: 0.5em 0 -0.3em 0;
 }
@@ -86,7 +142,14 @@ input[type="range"]{
     padding-right: 0.6em;
 }
 
+#mask_bug_info {
+  text-align: center;
+  display: block;
+  margin-bottom: 0.5em;
+}
 
+#txt2img_negative_prompt, #img2img_negative_prompt{
+}
 
 .progressDiv{
   width: 100%;

webui-user.bat

@@ -0,0 +1,8 @@
+@echo off
+
+set PYTHON=
+set GIT=
+set VENV_DIR=
+set COMMANDLINE_ARGS=
+
+call webui.bat

webui-user.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+###########################################
+# Change the variables below to your need:#
+###########################################
+
+# Install directory without trailing slash
+install_dir="/home/$(whoami)"
+
+# Name of the subdirectory (defaults to stable-diffusion-webui)
+clone_dir="stable-diffusion-webui"
+
+# Commandline arguments for webui.py, for example: export COMMANDLINE_ARGS=(--medvram --opt-split-attention)
+export COMMANDLINE_ARGS=()
+
+# python3 executable
+python_cmd="python3"
+
+# git executable
+export GIT=""
+
+# python3 venv without trailing slash (defaults to ${install_dir}/${clone_dir}/venv)
+venv_dir="venv"
+
+# install command for torch
+export TORCH_COMMAND=(python3 -m pip install torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113)
+
+# Requirements file to use for stable-diffusion-webui
+export REQS_FILE=""
+
+# Fixed git repos
+export K_DIFFUSION_PACKAGE=""
+export GFPGAN_PACKAGE=""
+
+# Fixed git commits
+export STABLE_DIFFUSION_COMMIT_HASH=""
+export TAMING_TRANSFORMERS_COMMIT_HASH=""
+export CODEFORMER_COMMIT_HASH=""
+export BLIP_COMMIT_HASH=""
+
+###########################################

webui.bat

@@ -1,28 +1,17 @@
 @echo off
 
-set PYTHON=python
-set GIT=git
-set COMMANDLINE_ARGS=
-set VENV_DIR=venv
+if not defined PYTHON (set PYTHON=python)
+if not defined VENV_DIR (set VENV_DIR=venv)
 
 mkdir tmp 2>NUL
 
-set TORCH_COMMAND=pip install torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
-set REQS_FILE=requirements_versions.txt
-
 %PYTHON% -c "" >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :check_git
+if %ERRORLEVEL% == 0 goto :start_venv
 echo Couldn't launch python
 goto :show_stdout_stderr
 
-:check_git
-%GIT% --help >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :setup_venv
-echo Couldn't launch git
-goto :show_stdout_stderr
-
-:setup_venv
-if [%VENV_DIR%] == [] goto :skip_venv
+:start_venv
+if [%VENV_DIR%] == [-] goto :skip_venv
 
 dir %VENV_DIR%\Scripts\Python.exe >tmp/stdout.txt 2>tmp/stderr.txt
 if %ERRORLEVEL% == 0 goto :activate_venv
@@ -35,114 +24,14 @@ echo Unable to create venv in directory %VENV_DIR%
 goto :show_stdout_stderr
 
 :activate_venv
-set PYTHON=%~dp0%VENV_DIR%\Scripts\Python.exe
-%PYTHON% --version
+set PYTHON="%~dp0%VENV_DIR%\Scripts\Python.exe"
 echo venv %PYTHON%
-goto :install_torch
+goto :launch
 
 :skip_venv
-%PYTHON% --version
-
-:install_torch
-
-%PYTHON% -c "import torch" >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :check_gpu
-echo Installing torch...
-%PYTHON% -m %TORCH_COMMAND% >tmp/stdout.txt 2>tmp/stderr.txt
-
-if %ERRORLEVEL% == 0 goto :check_gpu
-echo Failed to install torch
-goto :show_stdout_stderr
-
-:check_gpu
-%PYTHON% -c "import torch; assert torch.cuda.is_available(), 'CUDA is not available'" >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :install_sd_reqs
-echo Torch is not able to use GPU
-goto :show_stdout_stderr
-
-:install_sd_reqs
-%PYTHON% -c "import transformers; import wheel" >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :install_k_diff
-echo Installing SD requirements...
-%PYTHON% -m pip install wheel transformers==4.19.2 diffusers invisible-watermark --prefer-binary >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :install_k_diff
-goto :show_stdout_stderr
-
-:install_k_diff
-%PYTHON% -c "import k_diffusion.sampling" >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :install_GFPGAN
-echo Installing K-Diffusion...
-%PYTHON% -m pip install git+https://github.com/crowsonkb/k-diffusion.git --prefer-binary --only-binary=psutil >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :install_GFPGAN
-goto :show_stdout_stderr
-
-
-:install_GFPGAN
-%PYTHON% -c "import gfpgan" >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :install_reqs
-echo Installing GFPGAN
-%PYTHON% -m pip install git+https://github.com/TencentARC/GFPGAN.git --prefer-binary >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :install_reqs
-goto :show_stdout_stderr
-
-:install_reqs
-%PYTHON% -c "import omegaconf" >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :make_dirs
-echo Installing requirements...
-%PYTHON% -m pip install -r %REQS_FILE% --prefer-binary >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :update_numpy
-goto :show_stdout_stderr
-
-:update_numpy
-%PYTHON% -m pip install -U numpy --prefer-binary >tmp/stdout.txt 2>tmp/stderr.txt
-
-:make_dirs
-mkdir repositories 2>NUL
-
-if exist repositories\stable-diffusion goto :clone_transformers
-echo Cloning Stable Difusion repository...
-%GIT% clone https://github.com/CompVis/stable-diffusion.git repositories\stable-diffusion >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :clone_transformers
-goto :show_stdout_stderr
-
-:clone_transformers
-if exist repositories\taming-transformers goto :clone_codeformer
-echo Cloning Taming Transforming repository...
-%GIT% clone https://github.com/CompVis/taming-transformers.git repositories\taming-transformers >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :clone_codeformer
-goto :show_stdout_stderr
-
-:clone_codeformer
-if exist repositories\CodeFormer goto :install_codeformer_reqs
-echo Cloning CodeFormer repository...
-%GIT% clone https://github.com/sczhou/CodeFormer.git repositories\CodeFormer >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :install_codeformer_reqs
-goto :show_stdout_stderr
-
-:install_codeformer_reqs
-%PYTHON% -c "import lpips" >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :check_model
-echo Installing requirements for CodeFormer...
-%PYTHON% -m pip install -r repositories\CodeFormer\requirements.txt --prefer-binary >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :check_model
-goto :show_stdout_stderr
-
-
-:check_model
-dir model.ckpt >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :check_gfpgan
-echo Stable Diffusion model not found: you need to place model.ckpt file into same directory as this file.
-goto :show_stdout_stderr
-
-:check_gfpgan
-dir GFPGANv1.3.pth >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :launch
-echo GFPGAN not found: you need to place GFPGANv1.3.pth file into same directory as this file.
-echo Face fixing feature will not work.
 
 :launch
-echo Launching webui.py...
-%PYTHON% webui.py %COMMANDLINE_ARGS%
+%PYTHON% launch.py
 pause
 exit /b
 

webui.py

@@ -4,9 +4,7 @@ import threading
 from modules.paths import script_path
 
 import torch
-import numpy as np
 from omegaconf import OmegaConf
-from PIL import Image
 
 import signal
 
@@ -15,16 +13,14 @@ from ldm.util import instantiate_from_config
 from modules.shared import opts, cmd_opts, state
 import modules.shared as shared
 import modules.ui
-from modules.ui import plaintext_to_html
 import modules.scripts
-import modules.processing as processing
 import modules.sd_hijack
 import modules.codeformer_model
 import modules.gfpgan_model
 import modules.face_restoration
 import modules.realesrgan_model as realesrgan
 import modules.esrgan_model as esrgan
-import modules.images as images
+import modules.extras
 import modules.lowvram
 import modules.txt2img
 import modules.img2img
@@ -37,12 +33,14 @@ shared.face_restorers.append(modules.face_restoration.FaceRestoration())
 esrgan.load_models(cmd_opts.esrgan_models_path)
 realesrgan.setup_realesrgan()
 
+
 def load_model_from_config(config, ckpt, verbose=False):
-    print(f"Loading model from {ckpt}")
+    print(f"Loading model [{shared.sd_model_hash}] from {ckpt}")
     pl_sd = torch.load(ckpt, map_location="cpu")
     if "global_step" in pl_sd:
         print(f"Global Step: {pl_sd['global_step']}")
     sd = pl_sd["state_dict"]
+
     model = instantiate_from_config(config.model)
     m, u = model.load_state_dict(sd, strict=False)
     if len(m) > 0 and verbose:
@@ -51,84 +49,11 @@ def load_model_from_config(config, ckpt, verbose=False):
     if len(u) > 0 and verbose:
         print("unexpected keys:")
         print(u)
-
+    if cmd_opts.opt_channelslast:
+        model = model.to(memory_format=torch.channels_last)
     model.eval()
     return model
 
-cached_images = {}
-
-
-def run_extras(image, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility):
-    processing.torch_gc()
-
-    image = image.convert("RGB")
-
-    outpath = opts.outdir_samples or opts.outdir_extras_samples
-
-    if gfpgan_visibility > 0:
-        restored_img = modules.gfpgan_model.gfpgan_fix_faces(np.array(image, dtype=np.uint8))
-        res = Image.fromarray(restored_img)
-
-        if gfpgan_visibility < 1.0:
-            res = Image.blend(image, res, gfpgan_visibility)
-
-        image = res
-
-    if codeformer_visibility > 0:
-        restored_img = modules.codeformer_model.codeformer.restore(np.array(image, dtype=np.uint8), w=codeformer_weight)
-        res = Image.fromarray(restored_img)
-
-        if codeformer_visibility < 1.0:
-            res = Image.blend(image, res, codeformer_visibility)
-
-        image = res
-
-    if upscaling_resize != 1.0:
-        def upscale(image, scaler_index, resize):
-            small = image.crop((image.width // 2, image.height // 2, image.width // 2 + 10, image.height // 2 + 10))
-            pixels = tuple(np.array(small).flatten().tolist())
-            key = (resize, scaler_index, image.width, image.height) + pixels
-
-            c = cached_images.get(key)
-            if c is None:
-                upscaler = shared.sd_upscalers[scaler_index]
-                c = upscaler.upscale(image, image.width * resize, image.height * resize)
-                cached_images[key] = c
-
-            return c
-
-        res = upscale(image, extras_upscaler_1, upscaling_resize)
-
-        if extras_upscaler_2 != 0 and extras_upscaler_2_visibility>0:
-            res2 = upscale(image, extras_upscaler_2, upscaling_resize)
-            res = Image.blend(res, res2, extras_upscaler_2_visibility)
-
-        image = res
-
-    while len(cached_images) > 2:
-        del cached_images[next(iter(cached_images.keys()))]
-
-    images.save_image(image, outpath, "", None, '', opts.samples_format, short_filename=True, no_prompt=True)
-
-    return image, '', ''
-
-
-def run_pnginfo(image):
-    info = ''
-    for key, text in image.info.items():
-        info += f"""
-<div>
-<p><b>{plaintext_to_html(str(key))}</b></p>
-<p>{plaintext_to_html(str(text))}</p>
-</div>
-""".strip()+"\n"
-
-    if len(info) == 0:
-        message = "Nothing found in the image."
-        info = f"<div><p>{message}<p></div>"
-
-    return '', '', info
-
 
 queue_lock = threading.Lock()
 
@@ -153,6 +78,8 @@ def wrap_gradio_gpu_call(func):
     return modules.ui.wrap_gradio_call(f)
 
 
+modules.scripts.load_scripts(os.path.join(script_path, "scripts"))
+
 try:
     # this silences the annoying "Some weights of the model checkpoint were not used when initializing..." message at start.
 
@@ -162,6 +89,14 @@ try:
 except Exception:
     pass
 
+with open(cmd_opts.ckpt, "rb") as file:
+    import hashlib
+    m = hashlib.sha256()
+
+    file.seek(0x100000)
+    m.update(file.read(0x10000))
+    shared.sd_model_hash = m.hexdigest()[0:8]
+
 sd_config = OmegaConf.load(cmd_opts.config)
 shared.sd_model = load_model_from_config(sd_config, cmd_opts.ckpt)
 shared.sd_model = (shared.sd_model if cmd_opts.no_half else shared.sd_model.half())
@@ -173,22 +108,30 @@ else:
 
 modules.sd_hijack.model_hijack.hijack(shared.sd_model)
 
-modules.scripts.load_scripts(os.path.join(script_path, "scripts"))
 
-if __name__ == "__main__":
+def webui():
     # make the program just exit at ctrl+c without waiting for anything
     def sigint_handler(sig, frame):
         print(f'Interrupted with signal {sig} in {frame}')
         os._exit(0)
 
-
     signal.signal(signal.SIGINT, sigint_handler)
 
     demo = modules.ui.create_ui(
         txt2img=wrap_gradio_gpu_call(modules.txt2img.txt2img),
         img2img=wrap_gradio_gpu_call(modules.img2img.img2img),
-        run_extras=wrap_gradio_gpu_call(run_extras),
-        run_pnginfo=run_pnginfo
+        run_extras=wrap_gradio_gpu_call(modules.extras.run_extras),
+        run_pnginfo=modules.extras.run_pnginfo
     )
 
-    demo.launch(share=cmd_opts.share, server_name="0.0.0.0" if cmd_opts.listen else None)
+    demo.launch(
+        share=cmd_opts.share,
+        server_name="0.0.0.0" if cmd_opts.listen else None,
+        server_port=cmd_opts.port,
+        debug=cmd_opts.gradio_debug,
+        auth=[tuple(cred.split(':')) for cred in cmd_opts.gradio_auth.strip('"').split(',')] if cmd_opts.gradio_auth else None,
+    )
+
+
+if __name__ == "__main__":
+    webui()

webui.sh

@@ -0,0 +1,139 @@
+#!/bin/bash
+#################################################
+# Please do not make any changes to this file,  #
+# change the variables in webui-user.sh instead #
+#################################################
+# Read variables from webui-user.sh
+# shellcheck source=/dev/null
+if [[ -f webui-user.sh ]]
+then
+    source ./webui-user.sh
+fi
+
+# Set defaults
+# Install directory without trailing slash
+if [[ -z "${install_dir}" ]]
+then
+    install_dir="/home/$(whoami)"
+fi
+
+# Name of the subdirectory (defaults to stable-diffusion-webui)
+if [[ -z "${clone_dir}" ]]
+then
+    clone_dir="stable-diffusion-webui"
+fi
+
+# python3 executable
+if [[ -z "${python_cmd}" ]]
+then
+    python_cmd="python3"
+fi
+
+# git executable
+if [[ -z "${GIT}" ]]
+then
+    export GIT="git"
+fi
+
+# python3 venv without trailing slash (defaults to ${install_dir}/${clone_dir}/venv)
+if [[ -z "${venv_dir}" ]]
+then
+    venv_dir="venv"
+fi
+
+# install command for torch
+if [[ -z "${TORCH_COMMAND}" ]]
+then
+    export TORCH_COMMAND=(python3 -m pip install torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113)
+fi
+
+# Do not reinstall existing pip packages on Debian/Ubuntu
+export PIP_IGNORE_INSTALLED=0
+
+# Pretty print
+delimiter="################################################################"
+
+printf "\n%s\n" "${delimiter}"
+printf "\e[1m\e[32mInstall script for stable-diffusion + Web UI\n"
+printf "\e[1m\e[34mTested on Debian 11 (Bullseye)\e[0m"
+printf "\n%s\n" "${delimiter}"
+
+# Do not run as root
+if [[ $(id -u) -eq 0 ]]
+then
+    printf "\n%s\n" "${delimiter}"
+    printf "\e[1m\e[31mERROR: This script must not be launched as root, aborting...\e[0m"
+    printf "\n%s\n" "${delimiter}"
+    exit 1
+else
+    printf "\n%s\n" "${delimiter}"
+    printf "Running on \e[1m\e[32m%s\e[0m user" "$(whoami)"
+    printf "\n%s\n" "${delimiter}"
+fi
+
+if [[ -d .git ]]
+then
+    printf "\n%s\n" "${delimiter}"
+    printf "Repo already cloned, using it as install directory"
+    printf "\n%s\n" "${delimiter}"
+    install_dir="${PWD}/../"
+    clone_dir="${PWD##*/}"
+fi
+
+# Check prequisites
+for preq in git python3
+do
+    if ! hash "${preq}" &>/dev/null
+    then
+        printf "\n%s\n" "${delimiter}"
+        printf "\e[1m\e[31mERROR: %s is not installed, aborting...\e[0m" "${preq}"
+        printf "\n%s\n" "${delimiter}"
+        exit 1
+    fi
+done
+
+if ! "${python_cmd}" -c "import venv" &>/dev/null
+then
+    printf "\n%s\n" "${delimiter}"
+    printf "\e[1m\e[31mERROR: python3-venv is not installed, aborting...\e[0m"
+    printf "\n%s\n" "${delimiter}"
+    exit 1
+fi
+
+printf "\n%s\n" "${delimiter}"
+printf "Clone or update stable-diffusion-webui"
+printf "\n%s\n" "${delimiter}"
+cd "${install_dir}"/ || { printf "\e[1m\e[31mERROR: Can't cd to %s/, aborting...\e[0m" "${install_dir}"; exit 1; }
+if [[ -d "${clone_dir}" ]]
+then
+    cd "${clone_dir}"/ || { printf "\e[1m\e[31mERROR: Can't cd to %s/%s/, aborting...\e[0m" "${install_dir}" "${clone_dir}"; exit 1; }
+    "${GIT}" pull
+else
+    "${GIT}" clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git "${clone_dir}"
+    cd "${clone_dir}"/ || { printf "\e[1m\e[31mERROR: Can't cd to %s/%s/, aborting...\e[0m" "${install_dir}" "${clone_dir}"; exit 1; }
+fi
+
+printf "\n%s\n" "${delimiter}"
+printf "Create and activate python venv"
+printf "\n%s\n" "${delimiter}"
+cd "${install_dir}"/"${clone_dir}"/ || { printf "\e[1m\e[31mERROR: Can't cd to %s/%s/, aborting...\e[0m" "${install_dir}" "${clone_dir}"; exit 1; }
+if [[ ! -d "${venv_dir}" ]]
+then
+    "${python_cmd}" -m venv "${venv_dir}"
+    first_launch=1
+fi
+# shellcheck source=/dev/null
+if [[ -f "${venv_dir}"/bin/activate ]]
+then
+    source "${venv_dir}"/bin/activate
+else
+    printf "\n%s\n" "${delimiter}"
+    printf "\e[1m\e[31mERROR: Cannot activate python venv, aborting...\e[0m"
+    printf "\n%s\n" "${delimiter}"
+    exit 1
+fi
+
+printf "\n%s\n" "${delimiter}"
+printf "Launching launch.py..."
+printf "\n%s\n" "${delimiter}"
+"${python_cmd}" launch.py