Compare commits
9
Commits
codeformer
..
new-ui
| Author | SHA1 | Date | |
|---|---|---|---|
 AUTOMATIC
|
4374b498f6 | ||
|
AUTOMATIC
|
1b0c8de40e | ||
|
AUTOMATIC
|
469c992f15 | ||
 deggua
|
35ac3a6634 | ||
|
deggua
|
aa74ac7ebe | ||
|
deggua
|
19ba176dd3 | ||
|
deggua
|
021ba36ac2 | ||
|
deggua
|
8b8a000769 | ||
|
deggua
|
67a4549bf5 |
-11
@@ -1,11 +0,0 @@
|
||||
__pycache__
|
||||
/ESRGAN
|
||||
/repositories
|
||||
/venv
|
||||
/tmp
|
||||
/model.ckpt
|
||||
/GFPGANv1.3.pth
|
||||
/ui-config.json
|
||||
/outputs
|
||||
/config.json
|
||||
/log
|
||||
@@ -1,202 +1,55 @@
|
||||
# Stable Diffusion web UI
|
||||
A browser interface based on Gradio library for Stable Diffusion.
|
||||
|
||||
Original script with Gradio UI was written by a kind anonymous user. This is a modification.
|
||||
|
||||
|
||||
|
||||
## Feature showcase
|
||||
|
||||
[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)
|
||||
- Outpainting
|
||||
- Inpainting
|
||||
- Prompt matrix
|
||||
- Stable Diffusion upscale
|
||||
- Attention
|
||||
- Loopback
|
||||
- X/Y plot
|
||||
- Textual Inversion
|
||||
- Extras tab with:
|
||||
- GFPGAN, neural network that fixes faces
|
||||
- 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
|
||||
- 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
|
||||
- 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
|
||||
|
||||
## 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.
|
||||
### Stable Diffusion
|
||||
|
||||
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.
|
||||
This script assumes that you already have main Stable Diffusion sutff installed, assumed to be in directory `/sd`.
|
||||
If you don't have it installed, follow the guide:
|
||||
|
||||
You need `model.ckpt`, Stable Diffusion model checkpoint, a big file containing the neural network weights. You
|
||||
can obtain it from the following places:
|
||||
- [official download](https://huggingface.co/CompVis/stable-diffusion-v-1-4-original)
|
||||
- [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
|
||||
- https://rentry.org/kretard
|
||||
|
||||
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).
|
||||
This repository's `webgui.py` is a replacement for `kdiff.py` from the guide.
|
||||
|
||||
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).
|
||||
Particularly, following files must exist:
|
||||
|
||||
### Automatic installation/launch
|
||||
- `/sd/configs/stable-diffusion/v1-inference.yaml`
|
||||
- `/sd/models/ldm/stable-diffusion-v1/model.ckpt`
|
||||
- `/sd/ldm/util.py`
|
||||
- `/sd/k_diffusion/__init__.py`
|
||||
|
||||
- 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.
|
||||
### GFPGAN
|
||||
|
||||
#### Troublehooting:
|
||||
If you want to use GFPGAN to improve generated faces, you need to install it separately.
|
||||
Follow instructions from https://github.com/TencentARC/GFPGAN, but when cloning it, do so into Stable Diffusion main directory, `/sd`.
|
||||
After that download [GFPGANv1.3.pth](https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth) and put it
|
||||
into the `/sd/GFPGAN/experiments/pretrained_models` directory. If you're getting troubles with GFPGAN support, follow instructions
|
||||
from the GFPGAN's repository until `inference_gfpgan.py` script works.
|
||||
|
||||
- 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).
|
||||
- if you feel you broke something and want to reinstall from scratch, delete directories: `venv`, `repositories`.
|
||||
The following files must exist:
|
||||
|
||||
### Google collab
|
||||
- `/sd/GFPGAN/inference_gfpgan.py`
|
||||
- `/sd/GFPGAN/experiments/pretrained_models/GFPGANv1.3.pth`
|
||||
|
||||
If you don't want or can't run locally, here is google collab that allows you to run the webui:
|
||||
If the GFPGAN directory does not exist, you will not get the option to use GFPGAN in the UI. If it does exist, you will either be able
|
||||
to use it, or there will be a message in console with an error related to GFPGAN.
|
||||
|
||||
https://colab.research.google.com/drive/1Iy-xW9t1-OQWhb0hNxueGij8phCyluOh
|
||||
### Web UI
|
||||
|
||||
### What options to use for low VRAM videocards?
|
||||
- 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`
|
||||
also but the effect will likely be barely noticeable.
|
||||
- Otherwise, do not use any of those.
|
||||
Run the script as:
|
||||
|
||||
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.
|
||||
`python webui.py`
|
||||
|
||||
### Running online
|
||||
When running the script, you must be in the main Stable Diffusion directory, `/sd`. If you cloned this repository into a subdirectory
|
||||
of `/sd`, say, the `stable-diffusion-webui` directory, you will run it as:
|
||||
|
||||
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 `--listen` to make the server listen to network connections. This will allow computers on local newtork
|
||||
to access the UI, and if you configure port forwarding, also computers on the internet.
|
||||
|
||||
### 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.
|
||||
|
||||
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.
|
||||
|
||||
### How to change UI defaults?
|
||||
|
||||
After running once, a `ui-config.json` file appears in webui directory:
|
||||
|
||||
```json
|
||||
{
|
||||
"txt2img/Sampling Steps/value": 20,
|
||||
"txt2img/Sampling Steps/minimum": 1,
|
||||
"txt2img/Sampling Steps/maximum": 150,
|
||||
"txt2img/Sampling Steps/step": 1,
|
||||
"txt2img/Batch count/value": 1,
|
||||
"txt2img/Batch count/minimum": 1,
|
||||
"txt2img/Batch count/maximum": 32,
|
||||
"txt2img/Batch count/step": 1,
|
||||
"txt2img/Batch size/value": 1,
|
||||
"txt2img/Batch size/minimum": 1,
|
||||
```
|
||||
|
||||
Edit values to your liking and the next time you launch the program they will be applied.
|
||||
|
||||
|
||||
### Manual instructions
|
||||
Alternatively, if you don't want to run webui.bat, here are instructions for installing
|
||||
everything by hand:
|
||||
|
||||
```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%
|
||||
|
||||
:: 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.
|
||||
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
|
||||
git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git
|
||||
|
||||
:: install requirements of web ui
|
||||
pip install -r stable-diffusion-webui/requirements.txt
|
||||
|
||||
:: update numpy to latest version
|
||||
pip install -U numpy
|
||||
|
||||
:: (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
|
||||
|
||||
:: (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
|
||||
```
|
||||
|
||||
After that the installation is finished.
|
||||
|
||||
Run the command to start web ui:
|
||||
|
||||
```
|
||||
python stable-diffusion-webui/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 stable-diffusion-webui/webui.py`
|
||||
|
||||
When launching, you may get a very long warning message related to some weights not being used. You may freely ignore it.
|
||||
After a while, you will get a message like this:
|
||||
|
||||
```
|
||||
@@ -205,13 +58,207 @@ Running on local URL: http://127.0.0.1:7860/
|
||||
|
||||
Open the URL in browser, and you are good to go.
|
||||
|
||||
## Features
|
||||
The script creates a web UI for Stable Diffusion's txt2img and img2img scripts. Following are features added
|
||||
that are not in original script.
|
||||
|
||||
### Extras tab
|
||||
Additional neural network image improvement methods unrelated to stable diffusion.
|
||||
|
||||
#### GFPGAN
|
||||
Lets you improve faces in pictures using the GFPGAN model. There is a checkbox in every tab to use GFPGAN at 100%, and
|
||||
also a separate tab that just allows you to use GFPGAN on any picture, with a slider that controls how strongthe effect is.
|
||||
|
||||
|
||||
|
||||
#### Real-ESRGAN
|
||||
Image upscaler. You can choose from multiple models by original author, and specify by how much the image should be upscaled.
|
||||
Requires `realesrgan` librarty:
|
||||
|
||||
```commandline
|
||||
pip install realesrgan
|
||||
```
|
||||
|
||||
### Sampling method selection
|
||||
Pick out of multiple sampling methods for txt2img:
|
||||
|
||||
|
||||
|
||||
### Prompt matrix
|
||||
Separate multiple prompts using the `|` character, and the system will produce an image for every combination of them.
|
||||
For example, if you use `a busy city street in a modern city|illustration|cinematic lighting` prompt, there are four combinations possible (first part of prompt is always kept):
|
||||
|
||||
- `a busy city street in a modern city`
|
||||
- `a busy city street in a modern city, illustration`
|
||||
- `a busy city street in a modern city, cinematic lighting`
|
||||
- `a busy city street in a modern city, illustration, cinematic lighting`
|
||||
|
||||
Four images will be produced, in this order, all with same seed and each with corresponding prompt:
|
||||
|
||||
|
||||
Another example, this time with 5 prompts and 16 variations:
|
||||
|
||||
|
||||
If you use this feature, batch count will be ignored, because the number of pictures to produce
|
||||
depends on your prompts, but batch size will still work (generating multiple pictures at the
|
||||
same time for a small speed boost).
|
||||
|
||||
### Flagging
|
||||
Click the Flag button under the output section, and generated images will be saved to `log/images` directory, and generation parameters
|
||||
will be appended to a csv file `log/log.csv` in the `/sd` directory.
|
||||
|
||||
> but every image is saved, why would I need this?
|
||||
|
||||
If you're like me, you experiment a lot with prompts and settings, and only few images are worth saving. You can
|
||||
just save them using right click in browser, but then you won't be able to reproduce them later because you will not
|
||||
know what exact prompt created the image. If you use the flag button, generation parameters will be written to csv file,
|
||||
and you can easily find parameters for an image by searching for its filename.
|
||||
|
||||
### Copy-paste generation parameters
|
||||
A text output provides generation parameters in an easy to copy-paste form for easy sharing.
|
||||
|
||||
|
||||
|
||||
If you generate multiple pictures, the displayed seed will be the seed of the first one.
|
||||
|
||||
### Correct seeds for batches
|
||||
If you use a seed of 1000 to generate two batches of two images each, four generated images will have seeds: `1000, 1001, 1002, 1003`.
|
||||
Previous versions of the UI would produce `1000, x, 1001, x`, where x is an image that can't be generated by any seed.
|
||||
|
||||
### Resizing
|
||||
There are three options for resizing input images in img2img mode:
|
||||
|
||||
- Just resize - simply resizes source image to target resolution, resulting in incorrect aspect ratio
|
||||
- Crop and resize - resize source image preserving aspect ratio so that entirety of target resolution is occupied by it, and crop parts that stick out
|
||||
- Resize and fill - resize source image preserving aspect ratio so that it entirely fits target resolution, and fill empty space by rows/columns from source image
|
||||
|
||||
Example:
|
||||
|
||||
|
||||
### Loading
|
||||
Gradio's loading graphic has a very negative effect on the processing speed of the neural network.
|
||||
My RTX 3090 makes images about 10% faster when the tab with gradio is not active. By default, the UI
|
||||
now hides loading progress animation and replaces it with static "Loading..." text, which achieves
|
||||
the same effect. Use the `--no-progressbar-hiding` commandline option to revert this and show loading animations.
|
||||
|
||||
### Prompt validation
|
||||
Stable Diffusion has a limit for input text length. If your prompt is too long, you will get a
|
||||
warning in the text output field, showing which parts of your text were truncated and ignored by the model.
|
||||
|
||||
### Loopback
|
||||
A checkbox for img2img allowing to automatically feed output image as input for the next batch. Equivalent to
|
||||
saving output image, and replacing input image with it. Batch count setting controls how many iterations of
|
||||
this you get.
|
||||
|
||||
Usually, when doing this, you would choose one of many images for the next iteration yourself, so the usefulness
|
||||
of this feature may be questionable, but I've managed to get some very nice outputs with it that I wasn't abble
|
||||
to get otherwise.
|
||||
|
||||
Example: (cherrypicked result; original picture by anon)
|
||||
|
||||
|
||||
|
||||
### Png info
|
||||
Adds information about generation parameters to PNG as a text chunk. You
|
||||
can view this information later using any software that supports viewing
|
||||
PNG chunk info, for example: https://www.nayuki.io/page/png-file-chunk-inspector
|
||||
|
||||
|
||||
|
||||
### Textual Inversion
|
||||
Allows you to use pretrained textual inversion embeddings.
|
||||
See original site for details: https://textual-inversion.github.io/.
|
||||
I used lstein's repo for training embdedding: https://github.com/lstein/stable-diffusion; if
|
||||
you want to train your own, I recommend following the guide on his site.
|
||||
|
||||
No additional libraries/repositories are required to use pretrained embeddings.
|
||||
|
||||
To make use of pretrained embeddings, create `embeddings` directory in the root dir of Stable
|
||||
Diffusion and put your embeddings into it. They must be .pt files about 5Kb in size, each with only
|
||||
one trained embedding, and the filename (without .pt) will be the term you'd use in 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.
|
||||
|
||||
|
||||
|
||||
### Settings
|
||||
A tab with settings, allowing you to use UI to edit more than half of parameters that previously
|
||||
were commandline. Settings are saved to config.js file. Settings that remain as commandline
|
||||
options are ones that are required at startup.
|
||||
|
||||
### Attention
|
||||
Using `()` in prompt decreases model's attention to enclosed words, and `[]` increases it. You can combine
|
||||
multiple modifiers:
|
||||
|
||||
|
||||
|
||||
### SD upscale
|
||||
Upscale image using RealESRGAN and then go through tiles of the result, improving them with img2img.
|
||||
|
||||
Original idea by: https://github.com/jquesnelle/txt2imghd. This is an independent implementation.
|
||||
|
||||
To use this feature, tick a checkbox in the img2img interface. Original
|
||||
image will be upscaled to twice the original width and height, while width and height sliders
|
||||
will specify the size of individual tiles. At the moment this method does not support batch size.
|
||||
|
||||
|
||||
|
||||
### User scripts
|
||||
If the program is launched with `--allow-code` option, an extra text input field for script code
|
||||
is available in txt2img interface. It allows you to input python code that will do the work with
|
||||
image. If this field is not empty, the processing that would happen normally is skipped.
|
||||
|
||||
In code, access parameters from web UI using the `p` variable, and provide outputs for web UI
|
||||
using the `display(images, seed, info)` function. All globals from script are also accessible.
|
||||
|
||||
As an example, here is a script that draws a chart seen below (and also saves it as `test/gnomeplot/gnome5.png`):
|
||||
|
||||
```python
|
||||
steps = [4, 8,12,16, 20]
|
||||
cfg_scales = [5.0,10.0,15.0]
|
||||
|
||||
def cell(x, y, p=p):
|
||||
p.steps = x
|
||||
p.cfg_scale = y
|
||||
return process_images(p).images[0]
|
||||
|
||||
images = [draw_xy_grid(
|
||||
xs = steps,
|
||||
ys = cfg_scales,
|
||||
x_label = lambda x: f'Steps = {x}',
|
||||
y_label = lambda y: f'CFG = {y}',
|
||||
cell = cell
|
||||
)]
|
||||
|
||||
save_image(images[0], 'test/gnomeplot', 'gnome5')
|
||||
display(images)
|
||||
```
|
||||
|
||||
|
||||
|
||||
A more simple script that would just process the image and output it normally:
|
||||
|
||||
```python
|
||||
processed = process_images(p)
|
||||
|
||||
print("Seed was: " + str(processed.seed))
|
||||
|
||||
display(processed.images, processed.seed, processed.info)
|
||||
```
|
||||
|
||||
### `--lowvram`
|
||||
Optimizations for GPUs with low VRAM. This should make it possible to generate 512x512 images on videocards with 4GB memory.
|
||||
|
||||
The original idea of those optimizations is by basujindal: https://github.com/basujindal/stable-diffusion. Model is separated into modules,
|
||||
and only one module is kept in GPU memory; when another module needs to run, the previous is removed from GPU memory.
|
||||
|
||||
It should be obvious but the nature of those optimizations makes the processing run slower -- about 10 times slower
|
||||
compared to normal operation on my RTX 3090.
|
||||
|
||||
This is an independent implementation that does not require any modification to original Stable Diffusion code, and
|
||||
with all code concenrated in one place rather than scattered around the program.
|
||||
|
||||
|
||||
## 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
|
||||
- Idea for SD upscale - https://github.com/jquesnelle/txt2imghd
|
||||
- Initial Gradio script - posted on 4chan by an Anonymous user. Thank you Anonymous user.
|
||||
- (You)
|
||||
|
||||
-3051
File diff suppressed because it is too large
Load Diff
Binary file not shown.
|
After Width: | Height: | Size: 677 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 944 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 678 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 460 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 465 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 113 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.8 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.2 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 453 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 5.3 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.0 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 712 KiB |
@@ -1,25 +0,0 @@
|
||||
import os.path
|
||||
import csv
|
||||
from collections import namedtuple
|
||||
|
||||
Artist = namedtuple("Artist", ['name', 'weight', 'category'])
|
||||
|
||||
|
||||
class ArtistsDatabase:
|
||||
def __init__(self, filename):
|
||||
self.cats = set()
|
||||
self.artists = []
|
||||
|
||||
if not os.path.exists(filename):
|
||||
return
|
||||
|
||||
with open(filename, "r", newline='', encoding="utf8") as file:
|
||||
reader = csv.DictReader(file)
|
||||
|
||||
for row in reader:
|
||||
artist = Artist(row["artist"], float(row["score"]), row["category"])
|
||||
self.artists.append(artist)
|
||||
self.cats.add(artist.category)
|
||||
|
||||
def categories(self):
|
||||
return sorted(self.cats)
|
||||
@@ -1,278 +0,0 @@
|
||||
# this file is copied from CodeFormer repository. Please see comment in modules/codeformer_model.py
|
||||
|
||||
import math
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn, Tensor
|
||||
import torch.nn.functional as F
|
||||
from typing import Optional, List
|
||||
|
||||
from modules.codeformer.vqgan_arch import *
|
||||
from basicsr.utils import get_root_logger
|
||||
from basicsr.utils.registry import ARCH_REGISTRY
|
||||
|
||||
def calc_mean_std(feat, eps=1e-5):
|
||||
"""Calculate mean and std for adaptive_instance_normalization.
|
||||
|
||||
Args:
|
||||
feat (Tensor): 4D tensor.
|
||||
eps (float): A small value added to the variance to avoid
|
||||
divide-by-zero. Default: 1e-5.
|
||||
"""
|
||||
size = feat.size()
|
||||
assert len(size) == 4, 'The input feature should be 4D tensor.'
|
||||
b, c = size[:2]
|
||||
feat_var = feat.view(b, c, -1).var(dim=2) + eps
|
||||
feat_std = feat_var.sqrt().view(b, c, 1, 1)
|
||||
feat_mean = feat.view(b, c, -1).mean(dim=2).view(b, c, 1, 1)
|
||||
return feat_mean, feat_std
|
||||
|
||||
|
||||
def adaptive_instance_normalization(content_feat, style_feat):
|
||||
"""Adaptive instance normalization.
|
||||
|
||||
Adjust the reference features to have the similar color and illuminations
|
||||
as those in the degradate features.
|
||||
|
||||
Args:
|
||||
content_feat (Tensor): The reference feature.
|
||||
style_feat (Tensor): The degradate features.
|
||||
"""
|
||||
size = content_feat.size()
|
||||
style_mean, style_std = calc_mean_std(style_feat)
|
||||
content_mean, content_std = calc_mean_std(content_feat)
|
||||
normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)
|
||||
return normalized_feat * style_std.expand(size) + style_mean.expand(size)
|
||||
|
||||
|
||||
class PositionEmbeddingSine(nn.Module):
|
||||
"""
|
||||
This is a more standard version of the position embedding, very similar to the one
|
||||
used by the Attention is all you need paper, generalized to work on images.
|
||||
"""
|
||||
|
||||
def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
|
||||
super().__init__()
|
||||
self.num_pos_feats = num_pos_feats
|
||||
self.temperature = temperature
|
||||
self.normalize = normalize
|
||||
if scale is not None and normalize is False:
|
||||
raise ValueError("normalize should be True if scale is passed")
|
||||
if scale is None:
|
||||
scale = 2 * math.pi
|
||||
self.scale = scale
|
||||
|
||||
def forward(self, x, mask=None):
|
||||
if mask is None:
|
||||
mask = torch.zeros((x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool)
|
||||
not_mask = ~mask
|
||||
y_embed = not_mask.cumsum(1, dtype=torch.float32)
|
||||
x_embed = not_mask.cumsum(2, dtype=torch.float32)
|
||||
if self.normalize:
|
||||
eps = 1e-6
|
||||
y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
|
||||
x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
|
||||
|
||||
dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
|
||||
dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
|
||||
|
||||
pos_x = x_embed[:, :, :, None] / dim_t
|
||||
pos_y = y_embed[:, :, :, None] / dim_t
|
||||
pos_x = torch.stack(
|
||||
(pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4
|
||||
).flatten(3)
|
||||
pos_y = torch.stack(
|
||||
(pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4
|
||||
).flatten(3)
|
||||
pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
|
||||
return pos
|
||||
|
||||
def _get_activation_fn(activation):
|
||||
"""Return an activation function given a string"""
|
||||
if activation == "relu":
|
||||
return F.relu
|
||||
if activation == "gelu":
|
||||
return F.gelu
|
||||
if activation == "glu":
|
||||
return F.glu
|
||||
raise RuntimeError(F"activation should be relu/gelu, not {activation}.")
|
||||
|
||||
|
||||
class TransformerSALayer(nn.Module):
|
||||
def __init__(self, embed_dim, nhead=8, dim_mlp=2048, dropout=0.0, activation="gelu"):
|
||||
super().__init__()
|
||||
self.self_attn = nn.MultiheadAttention(embed_dim, nhead, dropout=dropout)
|
||||
# Implementation of Feedforward model - MLP
|
||||
self.linear1 = nn.Linear(embed_dim, dim_mlp)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.linear2 = nn.Linear(dim_mlp, embed_dim)
|
||||
|
||||
self.norm1 = nn.LayerNorm(embed_dim)
|
||||
self.norm2 = nn.LayerNorm(embed_dim)
|
||||
self.dropout1 = nn.Dropout(dropout)
|
||||
self.dropout2 = nn.Dropout(dropout)
|
||||
|
||||
self.activation = _get_activation_fn(activation)
|
||||
|
||||
def with_pos_embed(self, tensor, pos: Optional[Tensor]):
|
||||
return tensor if pos is None else tensor + pos
|
||||
|
||||
def forward(self, tgt,
|
||||
tgt_mask: Optional[Tensor] = None,
|
||||
tgt_key_padding_mask: Optional[Tensor] = None,
|
||||
query_pos: Optional[Tensor] = None):
|
||||
|
||||
# self attention
|
||||
tgt2 = self.norm1(tgt)
|
||||
q = k = self.with_pos_embed(tgt2, query_pos)
|
||||
tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
|
||||
key_padding_mask=tgt_key_padding_mask)[0]
|
||||
tgt = tgt + self.dropout1(tgt2)
|
||||
|
||||
# ffn
|
||||
tgt2 = self.norm2(tgt)
|
||||
tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
|
||||
tgt = tgt + self.dropout2(tgt2)
|
||||
return tgt
|
||||
|
||||
class Fuse_sft_block(nn.Module):
|
||||
def __init__(self, in_ch, out_ch):
|
||||
super().__init__()
|
||||
self.encode_enc = ResBlock(2*in_ch, out_ch)
|
||||
|
||||
self.scale = nn.Sequential(
|
||||
nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(0.2, True),
|
||||
nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1))
|
||||
|
||||
self.shift = nn.Sequential(
|
||||
nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(0.2, True),
|
||||
nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1))
|
||||
|
||||
def forward(self, enc_feat, dec_feat, w=1):
|
||||
enc_feat = self.encode_enc(torch.cat([enc_feat, dec_feat], dim=1))
|
||||
scale = self.scale(enc_feat)
|
||||
shift = self.shift(enc_feat)
|
||||
residual = w * (dec_feat * scale + shift)
|
||||
out = dec_feat + residual
|
||||
return out
|
||||
|
||||
|
||||
@ARCH_REGISTRY.register()
|
||||
class CodeFormer(VQAutoEncoder):
|
||||
def __init__(self, dim_embd=512, n_head=8, n_layers=9,
|
||||
codebook_size=1024, latent_size=256,
|
||||
connect_list=['32', '64', '128', '256'],
|
||||
fix_modules=['quantize','generator']):
|
||||
super(CodeFormer, self).__init__(512, 64, [1, 2, 2, 4, 4, 8], 'nearest',2, [16], codebook_size)
|
||||
|
||||
if fix_modules is not None:
|
||||
for module in fix_modules:
|
||||
for param in getattr(self, module).parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
self.connect_list = connect_list
|
||||
self.n_layers = n_layers
|
||||
self.dim_embd = dim_embd
|
||||
self.dim_mlp = dim_embd*2
|
||||
|
||||
self.position_emb = nn.Parameter(torch.zeros(latent_size, self.dim_embd))
|
||||
self.feat_emb = nn.Linear(256, self.dim_embd)
|
||||
|
||||
# transformer
|
||||
self.ft_layers = nn.Sequential(*[TransformerSALayer(embed_dim=dim_embd, nhead=n_head, dim_mlp=self.dim_mlp, dropout=0.0)
|
||||
for _ in range(self.n_layers)])
|
||||
|
||||
# logits_predict head
|
||||
self.idx_pred_layer = nn.Sequential(
|
||||
nn.LayerNorm(dim_embd),
|
||||
nn.Linear(dim_embd, codebook_size, bias=False))
|
||||
|
||||
self.channels = {
|
||||
'16': 512,
|
||||
'32': 256,
|
||||
'64': 256,
|
||||
'128': 128,
|
||||
'256': 128,
|
||||
'512': 64,
|
||||
}
|
||||
|
||||
# after second residual block for > 16, before attn layer for ==16
|
||||
self.fuse_encoder_block = {'512':2, '256':5, '128':8, '64':11, '32':14, '16':18}
|
||||
# after first residual block for > 16, before attn layer for ==16
|
||||
self.fuse_generator_block = {'16':6, '32': 9, '64':12, '128':15, '256':18, '512':21}
|
||||
|
||||
# fuse_convs_dict
|
||||
self.fuse_convs_dict = nn.ModuleDict()
|
||||
for f_size in self.connect_list:
|
||||
in_ch = self.channels[f_size]
|
||||
self.fuse_convs_dict[f_size] = Fuse_sft_block(in_ch, in_ch)
|
||||
|
||||
def _init_weights(self, module):
|
||||
if isinstance(module, (nn.Linear, nn.Embedding)):
|
||||
module.weight.data.normal_(mean=0.0, std=0.02)
|
||||
if isinstance(module, nn.Linear) and module.bias is not None:
|
||||
module.bias.data.zero_()
|
||||
elif isinstance(module, nn.LayerNorm):
|
||||
module.bias.data.zero_()
|
||||
module.weight.data.fill_(1.0)
|
||||
|
||||
def forward(self, x, w=0, detach_16=True, code_only=False, adain=False):
|
||||
# ################### Encoder #####################
|
||||
enc_feat_dict = {}
|
||||
out_list = [self.fuse_encoder_block[f_size] for f_size in self.connect_list]
|
||||
for i, block in enumerate(self.encoder.blocks):
|
||||
x = block(x)
|
||||
if i in out_list:
|
||||
enc_feat_dict[str(x.shape[-1])] = x.clone()
|
||||
|
||||
lq_feat = x
|
||||
# ################# Transformer ###################
|
||||
# quant_feat, codebook_loss, quant_stats = self.quantize(lq_feat)
|
||||
pos_emb = self.position_emb.unsqueeze(1).repeat(1,x.shape[0],1)
|
||||
# BCHW -> BC(HW) -> (HW)BC
|
||||
feat_emb = self.feat_emb(lq_feat.flatten(2).permute(2,0,1))
|
||||
query_emb = feat_emb
|
||||
# Transformer encoder
|
||||
for layer in self.ft_layers:
|
||||
query_emb = layer(query_emb, query_pos=pos_emb)
|
||||
|
||||
# output logits
|
||||
logits = self.idx_pred_layer(query_emb) # (hw)bn
|
||||
logits = logits.permute(1,0,2) # (hw)bn -> b(hw)n
|
||||
|
||||
if code_only: # for training stage II
|
||||
# logits doesn't need softmax before cross_entropy loss
|
||||
return logits, lq_feat
|
||||
|
||||
# ################# Quantization ###################
|
||||
# if self.training:
|
||||
# quant_feat = torch.einsum('btn,nc->btc', [soft_one_hot, self.quantize.embedding.weight])
|
||||
# # b(hw)c -> bc(hw) -> bchw
|
||||
# quant_feat = quant_feat.permute(0,2,1).view(lq_feat.shape)
|
||||
# ------------
|
||||
soft_one_hot = F.softmax(logits, dim=2)
|
||||
_, top_idx = torch.topk(soft_one_hot, 1, dim=2)
|
||||
quant_feat = self.quantize.get_codebook_feat(top_idx, shape=[x.shape[0],16,16,256])
|
||||
# preserve gradients
|
||||
# quant_feat = lq_feat + (quant_feat - lq_feat).detach()
|
||||
|
||||
if detach_16:
|
||||
quant_feat = quant_feat.detach() # for training stage III
|
||||
if adain:
|
||||
quant_feat = adaptive_instance_normalization(quant_feat, lq_feat)
|
||||
|
||||
# ################## Generator ####################
|
||||
x = quant_feat
|
||||
fuse_list = [self.fuse_generator_block[f_size] for f_size in self.connect_list]
|
||||
|
||||
for i, block in enumerate(self.generator.blocks):
|
||||
x = block(x)
|
||||
if i in fuse_list: # fuse after i-th block
|
||||
f_size = str(x.shape[-1])
|
||||
if w>0:
|
||||
x = self.fuse_convs_dict[f_size](enc_feat_dict[f_size].detach(), x, w)
|
||||
out = x
|
||||
# logits doesn't need softmax before cross_entropy loss
|
||||
return out, logits, lq_feat
|
||||
@@ -1,437 +0,0 @@
|
||||
# this file is copied from CodeFormer repository. Please see comment in modules/codeformer_model.py
|
||||
|
||||
'''
|
||||
VQGAN code, adapted from the original created by the Unleashing Transformers authors:
|
||||
https://github.com/samb-t/unleashing-transformers/blob/master/models/vqgan.py
|
||||
|
||||
'''
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import copy
|
||||
from basicsr.utils import get_root_logger
|
||||
from basicsr.utils.registry import ARCH_REGISTRY
|
||||
|
||||
def normalize(in_channels):
|
||||
return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def swish(x):
|
||||
return x*torch.sigmoid(x)
|
||||
|
||||
|
||||
# Define VQVAE classes
|
||||
class VectorQuantizer(nn.Module):
|
||||
def __init__(self, codebook_size, emb_dim, beta):
|
||||
super(VectorQuantizer, self).__init__()
|
||||
self.codebook_size = codebook_size # number of embeddings
|
||||
self.emb_dim = emb_dim # dimension of embedding
|
||||
self.beta = beta # commitment cost used in loss term, beta * ||z_e(x)-sg[e]||^2
|
||||
self.embedding = nn.Embedding(self.codebook_size, self.emb_dim)
|
||||
self.embedding.weight.data.uniform_(-1.0 / self.codebook_size, 1.0 / self.codebook_size)
|
||||
|
||||
def forward(self, z):
|
||||
# reshape z -> (batch, height, width, channel) and flatten
|
||||
z = z.permute(0, 2, 3, 1).contiguous()
|
||||
z_flattened = z.view(-1, self.emb_dim)
|
||||
|
||||
# distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
|
||||
d = (z_flattened ** 2).sum(dim=1, keepdim=True) + (self.embedding.weight**2).sum(1) - \
|
||||
2 * torch.matmul(z_flattened, self.embedding.weight.t())
|
||||
|
||||
mean_distance = torch.mean(d)
|
||||
# find closest encodings
|
||||
# min_encoding_indices = torch.argmin(d, dim=1).unsqueeze(1)
|
||||
min_encoding_scores, min_encoding_indices = torch.topk(d, 1, dim=1, largest=False)
|
||||
# [0-1], higher score, higher confidence
|
||||
min_encoding_scores = torch.exp(-min_encoding_scores/10)
|
||||
|
||||
min_encodings = torch.zeros(min_encoding_indices.shape[0], self.codebook_size).to(z)
|
||||
min_encodings.scatter_(1, min_encoding_indices, 1)
|
||||
|
||||
# get quantized latent vectors
|
||||
z_q = torch.matmul(min_encodings, self.embedding.weight).view(z.shape)
|
||||
# compute loss for embedding
|
||||
loss = torch.mean((z_q.detach()-z)**2) + self.beta * torch.mean((z_q - z.detach()) ** 2)
|
||||
# preserve gradients
|
||||
z_q = z + (z_q - z).detach()
|
||||
|
||||
# perplexity
|
||||
e_mean = torch.mean(min_encodings, dim=0)
|
||||
perplexity = torch.exp(-torch.sum(e_mean * torch.log(e_mean + 1e-10)))
|
||||
# reshape back to match original input shape
|
||||
z_q = z_q.permute(0, 3, 1, 2).contiguous()
|
||||
|
||||
return z_q, loss, {
|
||||
"perplexity": perplexity,
|
||||
"min_encodings": min_encodings,
|
||||
"min_encoding_indices": min_encoding_indices,
|
||||
"min_encoding_scores": min_encoding_scores,
|
||||
"mean_distance": mean_distance
|
||||
}
|
||||
|
||||
def get_codebook_feat(self, indices, shape):
|
||||
# input indices: batch*token_num -> (batch*token_num)*1
|
||||
# shape: batch, height, width, channel
|
||||
indices = indices.view(-1,1)
|
||||
min_encodings = torch.zeros(indices.shape[0], self.codebook_size).to(indices)
|
||||
min_encodings.scatter_(1, indices, 1)
|
||||
# get quantized latent vectors
|
||||
z_q = torch.matmul(min_encodings.float(), self.embedding.weight)
|
||||
|
||||
if shape is not None: # reshape back to match original input shape
|
||||
z_q = z_q.view(shape).permute(0, 3, 1, 2).contiguous()
|
||||
|
||||
return z_q
|
||||
|
||||
|
||||
class GumbelQuantizer(nn.Module):
|
||||
def __init__(self, codebook_size, emb_dim, num_hiddens, straight_through=False, kl_weight=5e-4, temp_init=1.0):
|
||||
super().__init__()
|
||||
self.codebook_size = codebook_size # number of embeddings
|
||||
self.emb_dim = emb_dim # dimension of embedding
|
||||
self.straight_through = straight_through
|
||||
self.temperature = temp_init
|
||||
self.kl_weight = kl_weight
|
||||
self.proj = nn.Conv2d(num_hiddens, codebook_size, 1) # projects last encoder layer to quantized logits
|
||||
self.embed = nn.Embedding(codebook_size, emb_dim)
|
||||
|
||||
def forward(self, z):
|
||||
hard = self.straight_through if self.training else True
|
||||
|
||||
logits = self.proj(z)
|
||||
|
||||
soft_one_hot = F.gumbel_softmax(logits, tau=self.temperature, dim=1, hard=hard)
|
||||
|
||||
z_q = torch.einsum("b n h w, n d -> b d h w", soft_one_hot, self.embed.weight)
|
||||
|
||||
# + kl divergence to the prior loss
|
||||
qy = F.softmax(logits, dim=1)
|
||||
diff = self.kl_weight * torch.sum(qy * torch.log(qy * self.codebook_size + 1e-10), dim=1).mean()
|
||||
min_encoding_indices = soft_one_hot.argmax(dim=1)
|
||||
|
||||
return z_q, diff, {
|
||||
"min_encoding_indices": min_encoding_indices
|
||||
}
|
||||
|
||||
|
||||
class Downsample(nn.Module):
|
||||
def __init__(self, in_channels):
|
||||
super().__init__()
|
||||
self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
|
||||
|
||||
def forward(self, x):
|
||||
pad = (0, 1, 0, 1)
|
||||
x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
|
||||
class Upsample(nn.Module):
|
||||
def __init__(self, in_channels):
|
||||
super().__init__()
|
||||
self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, x):
|
||||
x = F.interpolate(x, scale_factor=2.0, mode="nearest")
|
||||
x = self.conv(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class ResBlock(nn.Module):
|
||||
def __init__(self, in_channels, out_channels=None):
|
||||
super(ResBlock, self).__init__()
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = in_channels if out_channels is None else out_channels
|
||||
self.norm1 = normalize(in_channels)
|
||||
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
self.norm2 = normalize(out_channels)
|
||||
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
if self.in_channels != self.out_channels:
|
||||
self.conv_out = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
|
||||
|
||||
def forward(self, x_in):
|
||||
x = x_in
|
||||
x = self.norm1(x)
|
||||
x = swish(x)
|
||||
x = self.conv1(x)
|
||||
x = self.norm2(x)
|
||||
x = swish(x)
|
||||
x = self.conv2(x)
|
||||
if self.in_channels != self.out_channels:
|
||||
x_in = self.conv_out(x_in)
|
||||
|
||||
return x + x_in
|
||||
|
||||
|
||||
class AttnBlock(nn.Module):
|
||||
def __init__(self, in_channels):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
|
||||
self.norm = normalize(in_channels)
|
||||
self.q = torch.nn.Conv2d(
|
||||
in_channels,
|
||||
in_channels,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0
|
||||
)
|
||||
self.k = torch.nn.Conv2d(
|
||||
in_channels,
|
||||
in_channels,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0
|
||||
)
|
||||
self.v = torch.nn.Conv2d(
|
||||
in_channels,
|
||||
in_channels,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0
|
||||
)
|
||||
self.proj_out = torch.nn.Conv2d(
|
||||
in_channels,
|
||||
in_channels,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
h_ = x
|
||||
h_ = self.norm(h_)
|
||||
q = self.q(h_)
|
||||
k = self.k(h_)
|
||||
v = self.v(h_)
|
||||
|
||||
# compute attention
|
||||
b, c, h, w = q.shape
|
||||
q = q.reshape(b, c, h*w)
|
||||
q = q.permute(0, 2, 1)
|
||||
k = k.reshape(b, c, h*w)
|
||||
w_ = torch.bmm(q, k)
|
||||
w_ = w_ * (int(c)**(-0.5))
|
||||
w_ = F.softmax(w_, dim=2)
|
||||
|
||||
# attend to values
|
||||
v = v.reshape(b, c, h*w)
|
||||
w_ = w_.permute(0, 2, 1)
|
||||
h_ = torch.bmm(v, w_)
|
||||
h_ = h_.reshape(b, c, h, w)
|
||||
|
||||
h_ = self.proj_out(h_)
|
||||
|
||||
return x+h_
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
def __init__(self, in_channels, nf, emb_dim, ch_mult, num_res_blocks, resolution, attn_resolutions):
|
||||
super().__init__()
|
||||
self.nf = nf
|
||||
self.num_resolutions = len(ch_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.resolution = resolution
|
||||
self.attn_resolutions = attn_resolutions
|
||||
|
||||
curr_res = self.resolution
|
||||
in_ch_mult = (1,)+tuple(ch_mult)
|
||||
|
||||
blocks = []
|
||||
# initial convultion
|
||||
blocks.append(nn.Conv2d(in_channels, nf, kernel_size=3, stride=1, padding=1))
|
||||
|
||||
# residual and downsampling blocks, with attention on smaller res (16x16)
|
||||
for i in range(self.num_resolutions):
|
||||
block_in_ch = nf * in_ch_mult[i]
|
||||
block_out_ch = nf * ch_mult[i]
|
||||
for _ in range(self.num_res_blocks):
|
||||
blocks.append(ResBlock(block_in_ch, block_out_ch))
|
||||
block_in_ch = block_out_ch
|
||||
if curr_res in attn_resolutions:
|
||||
blocks.append(AttnBlock(block_in_ch))
|
||||
|
||||
if i != self.num_resolutions - 1:
|
||||
blocks.append(Downsample(block_in_ch))
|
||||
curr_res = curr_res // 2
|
||||
|
||||
# non-local attention block
|
||||
blocks.append(ResBlock(block_in_ch, block_in_ch))
|
||||
blocks.append(AttnBlock(block_in_ch))
|
||||
blocks.append(ResBlock(block_in_ch, block_in_ch))
|
||||
|
||||
# normalise and convert to latent size
|
||||
blocks.append(normalize(block_in_ch))
|
||||
blocks.append(nn.Conv2d(block_in_ch, emb_dim, kernel_size=3, stride=1, padding=1))
|
||||
self.blocks = nn.ModuleList(blocks)
|
||||
|
||||
def forward(self, x):
|
||||
for block in self.blocks:
|
||||
x = block(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class Generator(nn.Module):
|
||||
def __init__(self, nf, emb_dim, ch_mult, res_blocks, img_size, attn_resolutions):
|
||||
super().__init__()
|
||||
self.nf = nf
|
||||
self.ch_mult = ch_mult
|
||||
self.num_resolutions = len(self.ch_mult)
|
||||
self.num_res_blocks = res_blocks
|
||||
self.resolution = img_size
|
||||
self.attn_resolutions = attn_resolutions
|
||||
self.in_channels = emb_dim
|
||||
self.out_channels = 3
|
||||
block_in_ch = self.nf * self.ch_mult[-1]
|
||||
curr_res = self.resolution // 2 ** (self.num_resolutions-1)
|
||||
|
||||
blocks = []
|
||||
# initial conv
|
||||
blocks.append(nn.Conv2d(self.in_channels, block_in_ch, kernel_size=3, stride=1, padding=1))
|
||||
|
||||
# non-local attention block
|
||||
blocks.append(ResBlock(block_in_ch, block_in_ch))
|
||||
blocks.append(AttnBlock(block_in_ch))
|
||||
blocks.append(ResBlock(block_in_ch, block_in_ch))
|
||||
|
||||
for i in reversed(range(self.num_resolutions)):
|
||||
block_out_ch = self.nf * self.ch_mult[i]
|
||||
|
||||
for _ in range(self.num_res_blocks):
|
||||
blocks.append(ResBlock(block_in_ch, block_out_ch))
|
||||
block_in_ch = block_out_ch
|
||||
|
||||
if curr_res in self.attn_resolutions:
|
||||
blocks.append(AttnBlock(block_in_ch))
|
||||
|
||||
if i != 0:
|
||||
blocks.append(Upsample(block_in_ch))
|
||||
curr_res = curr_res * 2
|
||||
|
||||
blocks.append(normalize(block_in_ch))
|
||||
blocks.append(nn.Conv2d(block_in_ch, self.out_channels, kernel_size=3, stride=1, padding=1))
|
||||
|
||||
self.blocks = nn.ModuleList(blocks)
|
||||
|
||||
|
||||
def forward(self, x):
|
||||
for block in self.blocks:
|
||||
x = block(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
@ARCH_REGISTRY.register()
|
||||
class VQAutoEncoder(nn.Module):
|
||||
def __init__(self, img_size, nf, ch_mult, quantizer="nearest", res_blocks=2, attn_resolutions=[16], codebook_size=1024, emb_dim=256,
|
||||
beta=0.25, gumbel_straight_through=False, gumbel_kl_weight=1e-8, model_path=None):
|
||||
super().__init__()
|
||||
logger = get_root_logger()
|
||||
self.in_channels = 3
|
||||
self.nf = nf
|
||||
self.n_blocks = res_blocks
|
||||
self.codebook_size = codebook_size
|
||||
self.embed_dim = emb_dim
|
||||
self.ch_mult = ch_mult
|
||||
self.resolution = img_size
|
||||
self.attn_resolutions = attn_resolutions
|
||||
self.quantizer_type = quantizer
|
||||
self.encoder = Encoder(
|
||||
self.in_channels,
|
||||
self.nf,
|
||||
self.embed_dim,
|
||||
self.ch_mult,
|
||||
self.n_blocks,
|
||||
self.resolution,
|
||||
self.attn_resolutions
|
||||
)
|
||||
if self.quantizer_type == "nearest":
|
||||
self.beta = beta #0.25
|
||||
self.quantize = VectorQuantizer(self.codebook_size, self.embed_dim, self.beta)
|
||||
elif self.quantizer_type == "gumbel":
|
||||
self.gumbel_num_hiddens = emb_dim
|
||||
self.straight_through = gumbel_straight_through
|
||||
self.kl_weight = gumbel_kl_weight
|
||||
self.quantize = GumbelQuantizer(
|
||||
self.codebook_size,
|
||||
self.embed_dim,
|
||||
self.gumbel_num_hiddens,
|
||||
self.straight_through,
|
||||
self.kl_weight
|
||||
)
|
||||
self.generator = Generator(
|
||||
self.nf,
|
||||
self.embed_dim,
|
||||
self.ch_mult,
|
||||
self.n_blocks,
|
||||
self.resolution,
|
||||
self.attn_resolutions
|
||||
)
|
||||
|
||||
if model_path is not None:
|
||||
chkpt = torch.load(model_path, map_location='cpu')
|
||||
if 'params_ema' in chkpt:
|
||||
self.load_state_dict(torch.load(model_path, map_location='cpu')['params_ema'])
|
||||
logger.info(f'vqgan is loaded from: {model_path} [params_ema]')
|
||||
elif 'params' in chkpt:
|
||||
self.load_state_dict(torch.load(model_path, map_location='cpu')['params'])
|
||||
logger.info(f'vqgan is loaded from: {model_path} [params]')
|
||||
else:
|
||||
raise ValueError(f'Wrong params!')
|
||||
|
||||
|
||||
def forward(self, x):
|
||||
x = self.encoder(x)
|
||||
quant, codebook_loss, quant_stats = self.quantize(x)
|
||||
x = self.generator(quant)
|
||||
return x, codebook_loss, quant_stats
|
||||
|
||||
|
||||
|
||||
# patch based discriminator
|
||||
@ARCH_REGISTRY.register()
|
||||
class VQGANDiscriminator(nn.Module):
|
||||
def __init__(self, nc=3, ndf=64, n_layers=4, model_path=None):
|
||||
super().__init__()
|
||||
|
||||
layers = [nn.Conv2d(nc, ndf, kernel_size=4, stride=2, padding=1), nn.LeakyReLU(0.2, True)]
|
||||
ndf_mult = 1
|
||||
ndf_mult_prev = 1
|
||||
for n in range(1, n_layers): # gradually increase the number of filters
|
||||
ndf_mult_prev = ndf_mult
|
||||
ndf_mult = min(2 ** n, 8)
|
||||
layers += [
|
||||
nn.Conv2d(ndf * ndf_mult_prev, ndf * ndf_mult, kernel_size=4, stride=2, padding=1, bias=False),
|
||||
nn.BatchNorm2d(ndf * ndf_mult),
|
||||
nn.LeakyReLU(0.2, True)
|
||||
]
|
||||
|
||||
ndf_mult_prev = ndf_mult
|
||||
ndf_mult = min(2 ** n_layers, 8)
|
||||
|
||||
layers += [
|
||||
nn.Conv2d(ndf * ndf_mult_prev, ndf * ndf_mult, kernel_size=4, stride=1, padding=1, bias=False),
|
||||
nn.BatchNorm2d(ndf * ndf_mult),
|
||||
nn.LeakyReLU(0.2, True)
|
||||
]
|
||||
|
||||
layers += [
|
||||
nn.Conv2d(ndf * ndf_mult, 1, kernel_size=4, stride=1, padding=1)] # output 1 channel prediction map
|
||||
self.main = nn.Sequential(*layers)
|
||||
|
||||
if model_path is not None:
|
||||
chkpt = torch.load(model_path, map_location='cpu')
|
||||
if 'params_d' in chkpt:
|
||||
self.load_state_dict(torch.load(model_path, map_location='cpu')['params_d'])
|
||||
elif 'params' in chkpt:
|
||||
self.load_state_dict(torch.load(model_path, map_location='cpu')['params'])
|
||||
else:
|
||||
raise ValueError(f'Wrong params!')
|
||||
|
||||
def forward(self, x):
|
||||
return self.main(x)
|
||||
@@ -1,112 +0,0 @@
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
import torch
|
||||
|
||||
from modules import shared
|
||||
from modules.paths import script_path
|
||||
import modules.shared
|
||||
import modules.face_restoration
|
||||
from importlib import reload
|
||||
|
||||
# codeformer people made a choice to include modified basicsr librry to their projectwhich makes
|
||||
# it utterly impossiblr to use it alongside with other libraries that also use basicsr, like GFPGAN.
|
||||
# I am making a choice to include some files from codeformer to work around this issue.
|
||||
|
||||
pretrain_model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
|
||||
|
||||
have_codeformer = False
|
||||
codeformer = None
|
||||
|
||||
def setup_codeformer():
|
||||
path = modules.paths.paths.get("CodeFormer", None)
|
||||
if path is None:
|
||||
return
|
||||
|
||||
|
||||
# both GFPGAN and CodeFormer use bascisr, one has it installed from pip the other uses its own
|
||||
#stored_sys_path = sys.path
|
||||
#sys.path = [path] + sys.path
|
||||
|
||||
try:
|
||||
from torchvision.transforms.functional import normalize
|
||||
from modules.codeformer.codeformer_arch import CodeFormer
|
||||
from basicsr.utils.download_util import load_file_from_url
|
||||
from basicsr.utils import imwrite, img2tensor, tensor2img
|
||||
from facelib.utils.face_restoration_helper import FaceRestoreHelper
|
||||
from modules.shared import cmd_opts
|
||||
|
||||
net_class = CodeFormer
|
||||
|
||||
class FaceRestorerCodeFormer(modules.face_restoration.FaceRestoration):
|
||||
def name(self):
|
||||
return "CodeFormer"
|
||||
|
||||
def __init__(self):
|
||||
self.net = None
|
||||
self.face_helper = None
|
||||
|
||||
def create_models(self):
|
||||
|
||||
if self.net is not None and self.face_helper is not None:
|
||||
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)
|
||||
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)
|
||||
|
||||
if not cmd_opts.unload_gfpgan:
|
||||
self.net = net
|
||||
self.face_helper = face_helper
|
||||
|
||||
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()
|
||||
|
||||
for idx, cropped_face in enumerate(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)
|
||||
|
||||
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]
|
||||
restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
|
||||
del output
|
||||
torch.cuda.empty_cache()
|
||||
except Exception as error:
|
||||
print(f'\tFailed inference for CodeFormer: {error}', file=sys.stderr)
|
||||
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)
|
||||
|
||||
face_helper.get_inverse_affine(None)
|
||||
|
||||
restored_img = face_helper.paste_faces_to_input_image()
|
||||
restored_img = restored_img[:, :, ::-1]
|
||||
return restored_img
|
||||
|
||||
global have_codeformer
|
||||
have_codeformer = True
|
||||
|
||||
global codeformer
|
||||
codeformer = FaceRestorerCodeFormer()
|
||||
shared.face_restorers.append(codeformer)
|
||||
|
||||
except Exception:
|
||||
print("Error setting up CodeFormer:", file=sys.stderr)
|
||||
print(traceback.format_exc(), file=sys.stderr)
|
||||
|
||||
# sys.path = stored_sys_path
|
||||
@@ -1,80 +0,0 @@
|
||||
# this file is taken from https://github.com/xinntao/ESRGAN
|
||||
|
||||
import functools
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def make_layer(block, n_layers):
|
||||
layers = []
|
||||
for _ in range(n_layers):
|
||||
layers.append(block())
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
|
||||
class ResidualDenseBlock_5C(nn.Module):
|
||||
def __init__(self, nf=64, gc=32, bias=True):
|
||||
super(ResidualDenseBlock_5C, self).__init__()
|
||||
# gc: growth channel, i.e. intermediate channels
|
||||
self.conv1 = nn.Conv2d(nf, gc, 3, 1, 1, bias=bias)
|
||||
self.conv2 = nn.Conv2d(nf + gc, gc, 3, 1, 1, bias=bias)
|
||||
self.conv3 = nn.Conv2d(nf + 2 * gc, gc, 3, 1, 1, bias=bias)
|
||||
self.conv4 = nn.Conv2d(nf + 3 * gc, gc, 3, 1, 1, bias=bias)
|
||||
self.conv5 = nn.Conv2d(nf + 4 * gc, nf, 3, 1, 1, bias=bias)
|
||||
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
|
||||
|
||||
# initialization
|
||||
# mutil.initialize_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)
|
||||
|
||||
def forward(self, x):
|
||||
x1 = self.lrelu(self.conv1(x))
|
||||
x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
|
||||
x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
|
||||
x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
|
||||
x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
|
||||
return x5 * 0.2 + x
|
||||
|
||||
|
||||
class RRDB(nn.Module):
|
||||
'''Residual in Residual Dense Block'''
|
||||
|
||||
def __init__(self, nf, gc=32):
|
||||
super(RRDB, self).__init__()
|
||||
self.RDB1 = ResidualDenseBlock_5C(nf, gc)
|
||||
self.RDB2 = ResidualDenseBlock_5C(nf, gc)
|
||||
self.RDB3 = ResidualDenseBlock_5C(nf, gc)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.RDB1(x)
|
||||
out = self.RDB2(out)
|
||||
out = self.RDB3(out)
|
||||
return out * 0.2 + x
|
||||
|
||||
|
||||
class RRDBNet(nn.Module):
|
||||
def __init__(self, in_nc, out_nc, nf, nb, gc=32):
|
||||
super(RRDBNet, self).__init__()
|
||||
RRDB_block_f = functools.partial(RRDB, nf=nf, gc=gc)
|
||||
|
||||
self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
|
||||
self.RRDB_trunk = make_layer(RRDB_block_f, nb)
|
||||
self.trunk_conv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
|
||||
#### upsampling
|
||||
self.upconv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
|
||||
self.upconv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
|
||||
self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
|
||||
self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)
|
||||
|
||||
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
fea = self.conv_first(x)
|
||||
trunk = self.trunk_conv(self.RRDB_trunk(fea))
|
||||
fea = fea + trunk
|
||||
|
||||
fea = self.lrelu(self.upconv1(F.interpolate(fea, scale_factor=2, mode='nearest')))
|
||||
fea = self.lrelu(self.upconv2(F.interpolate(fea, scale_factor=2, mode='nearest')))
|
||||
out = self.conv_last(self.lrelu(self.HRconv(fea)))
|
||||
|
||||
return out
|
||||
@@ -1,134 +0,0 @@
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
import modules.esrgam_model_arch as arch
|
||||
from modules import shared
|
||||
from modules.shared import opts
|
||||
import modules.images
|
||||
|
||||
|
||||
def load_model(filename):
|
||||
# this code is adapted from https://github.com/xinntao/ESRGAN
|
||||
|
||||
pretrained_net = torch.load(filename)
|
||||
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
|
||||
|
||||
crt_net = crt_model.state_dict()
|
||||
load_net_clean = {}
|
||||
for k, v in pretrained_net.items():
|
||||
if k.startswith('module.'):
|
||||
load_net_clean[k[7:]] = v
|
||||
else:
|
||||
load_net_clean[k] = v
|
||||
pretrained_net = load_net_clean
|
||||
|
||||
tbd = []
|
||||
for k, v in crt_net.items():
|
||||
tbd.append(k)
|
||||
|
||||
# directly copy
|
||||
for k, v in crt_net.items():
|
||||
if k in pretrained_net and pretrained_net[k].size() == v.size():
|
||||
crt_net[k] = pretrained_net[k]
|
||||
tbd.remove(k)
|
||||
|
||||
crt_net['conv_first.weight'] = pretrained_net['model.0.weight']
|
||||
crt_net['conv_first.bias'] = pretrained_net['model.0.bias']
|
||||
|
||||
for k in tbd.copy():
|
||||
if 'RDB' in k:
|
||||
ori_k = k.replace('RRDB_trunk.', 'model.1.sub.')
|
||||
if '.weight' in k:
|
||||
ori_k = ori_k.replace('.weight', '.0.weight')
|
||||
elif '.bias' in k:
|
||||
ori_k = ori_k.replace('.bias', '.0.bias')
|
||||
crt_net[k] = pretrained_net[ori_k]
|
||||
tbd.remove(k)
|
||||
|
||||
crt_net['trunk_conv.weight'] = pretrained_net['model.1.sub.23.weight']
|
||||
crt_net['trunk_conv.bias'] = pretrained_net['model.1.sub.23.bias']
|
||||
crt_net['upconv1.weight'] = pretrained_net['model.3.weight']
|
||||
crt_net['upconv1.bias'] = pretrained_net['model.3.bias']
|
||||
crt_net['upconv2.weight'] = pretrained_net['model.6.weight']
|
||||
crt_net['upconv2.bias'] = pretrained_net['model.6.bias']
|
||||
crt_net['HRconv.weight'] = pretrained_net['model.8.weight']
|
||||
crt_net['HRconv.bias'] = pretrained_net['model.8.bias']
|
||||
crt_net['conv_last.weight'] = pretrained_net['model.10.weight']
|
||||
crt_net['conv_last.bias'] = pretrained_net['model.10.bias']
|
||||
|
||||
crt_model.load_state_dict(crt_net)
|
||||
crt_model.eval()
|
||||
return crt_model
|
||||
|
||||
def upscale_without_tiling(model, img):
|
||||
img = np.array(img)
|
||||
img = img[:, :, ::-1]
|
||||
img = np.moveaxis(img, 2, 0) / 255
|
||||
img = torch.from_numpy(img).float()
|
||||
img = img.unsqueeze(0).to(shared.device)
|
||||
with torch.no_grad():
|
||||
output = model(img)
|
||||
output = output.squeeze().float().cpu().clamp_(0, 1).numpy()
|
||||
output = 255. * np.moveaxis(output, 0, 2)
|
||||
output = output.astype(np.uint8)
|
||||
output = output[:, :, ::-1]
|
||||
return Image.fromarray(output, 'RGB')
|
||||
|
||||
|
||||
def esrgan_upscale(model, img):
|
||||
if opts.ESRGAN_tile == 0:
|
||||
return upscale_without_tiling(model, img)
|
||||
|
||||
grid = modules.images.split_grid(img, opts.ESRGAN_tile, opts.ESRGAN_tile, opts.ESRGAN_tile_overlap)
|
||||
newtiles = []
|
||||
scale_factor = 1
|
||||
|
||||
for y, h, row in grid.tiles:
|
||||
newrow = []
|
||||
for tiledata in row:
|
||||
x, w, tile = tiledata
|
||||
|
||||
output = upscale_without_tiling(model, tile)
|
||||
scale_factor = output.width // tile.width
|
||||
|
||||
newrow.append([x * scale_factor, w * scale_factor, output])
|
||||
newtiles.append([y * scale_factor, h * scale_factor, newrow])
|
||||
|
||||
newgrid = modules.images.Grid(newtiles, grid.tile_w * scale_factor, grid.tile_h * scale_factor, grid.image_w * scale_factor, grid.image_h * scale_factor, grid.overlap * scale_factor)
|
||||
output = modules.images.combine_grid(newgrid)
|
||||
return output
|
||||
|
||||
|
||||
class UpscalerESRGAN(modules.images.Upscaler):
|
||||
def __init__(self, filename, title):
|
||||
self.name = title
|
||||
self.model = load_model(filename)
|
||||
|
||||
def do_upscale(self, img):
|
||||
model = self.model.to(shared.device)
|
||||
img = esrgan_upscale(model, img)
|
||||
return img
|
||||
|
||||
|
||||
def load_models(dirname):
|
||||
for file in os.listdir(dirname):
|
||||
path = os.path.join(dirname, file)
|
||||
model_name, extension = os.path.splitext(file)
|
||||
|
||||
if extension != '.pt' and extension != '.pth':
|
||||
continue
|
||||
|
||||
try:
|
||||
modules.shared.sd_upscalers.append(UpscalerESRGAN(path, model_name))
|
||||
except Exception:
|
||||
print(f"Error loading ESRGAN model: {path}", file=sys.stderr)
|
||||
print(traceback.format_exc(), file=sys.stderr)
|
||||
@@ -1,19 +0,0 @@
|
||||
from modules import shared
|
||||
|
||||
|
||||
class FaceRestoration:
|
||||
def name(self):
|
||||
return "None"
|
||||
|
||||
def restore(self, np_image):
|
||||
return np_image
|
||||
|
||||
|
||||
def restore_faces(np_image):
|
||||
face_restorers = [x for x in shared.face_restorers if x.name() == shared.opts.face_restoration_model or shared.opts.face_restoration_model is None]
|
||||
if len(face_restorers) == 0:
|
||||
return np_image
|
||||
|
||||
face_restorer = face_restorers[0]
|
||||
|
||||
return face_restorer.restore(np_image)
|
||||
@@ -1,83 +0,0 @@
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
|
||||
from modules import shared
|
||||
from modules.shared import cmd_opts
|
||||
from modules.paths import script_path
|
||||
import modules.face_restoration
|
||||
|
||||
|
||||
def gfpgan_model_path():
|
||||
from modules.shared import cmd_opts
|
||||
|
||||
places = [script_path, '.', os.path.join(cmd_opts.gfpgan_dir, 'experiments/pretrained_models')]
|
||||
files = [cmd_opts.gfpgan_model] + [os.path.join(dirname, cmd_opts.gfpgan_model) for dirname in places]
|
||||
found = [x for x in files if os.path.exists(x)]
|
||||
|
||||
if len(found) == 0:
|
||||
raise Exception("GFPGAN model not found in paths: " + ", ".join(files))
|
||||
|
||||
return found[0]
|
||||
|
||||
|
||||
loaded_gfpgan_model = None
|
||||
|
||||
|
||||
def gfpgan():
|
||||
global loaded_gfpgan_model
|
||||
|
||||
if loaded_gfpgan_model is not None:
|
||||
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
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def gfpgan_fix_faces(np_image):
|
||||
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)
|
||||
np_image = gfpgan_output_bgr[:, :, ::-1]
|
||||
|
||||
return np_image
|
||||
|
||||
|
||||
have_gfpgan = False
|
||||
gfpgan_constructor = None
|
||||
|
||||
def setup_gfpgan():
|
||||
try:
|
||||
gfpgan_model_path()
|
||||
|
||||
if os.path.exists(cmd_opts.gfpgan_dir):
|
||||
sys.path.append(os.path.abspath(cmd_opts.gfpgan_dir))
|
||||
from gfpgan import GFPGANer
|
||||
|
||||
global have_gfpgan
|
||||
have_gfpgan = True
|
||||
|
||||
global gfpgan_constructor
|
||||
gfpgan_constructor = GFPGANer
|
||||
|
||||
class FaceRestorerGFPGAN(modules.face_restoration.FaceRestoration):
|
||||
def name(self):
|
||||
return "GFPGAN"
|
||||
|
||||
def restore(self, np_image):
|
||||
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)
|
||||
np_image = gfpgan_output_bgr[:, :, ::-1]
|
||||
|
||||
return np_image
|
||||
|
||||
shared.face_restorers.append(FaceRestorerGFPGAN())
|
||||
except Exception:
|
||||
print("Error setting up GFPGAN:", file=sys.stderr)
|
||||
print(traceback.format_exc(), file=sys.stderr)
|
||||
@@ -1,323 +0,0 @@
|
||||
import math
|
||||
import os
|
||||
from collections import namedtuple
|
||||
import re
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image, ImageFont, ImageDraw, PngImagePlugin
|
||||
|
||||
import modules.shared
|
||||
from modules.shared import opts
|
||||
|
||||
LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)
|
||||
|
||||
|
||||
def image_grid(imgs, batch_size=1, rows=None):
|
||||
if rows is None:
|
||||
if opts.n_rows > 0:
|
||||
rows = opts.n_rows
|
||||
elif opts.n_rows == 0:
|
||||
rows = batch_size
|
||||
else:
|
||||
rows = math.sqrt(len(imgs))
|
||||
rows = round(rows)
|
||||
|
||||
cols = math.ceil(len(imgs) / rows)
|
||||
|
||||
w, h = imgs[0].size
|
||||
grid = Image.new('RGB', size=(cols * w, rows * h), color='black')
|
||||
|
||||
for i, img in enumerate(imgs):
|
||||
grid.paste(img, box=(i % cols * w, i // cols * h))
|
||||
|
||||
return grid
|
||||
|
||||
|
||||
Grid = namedtuple("Grid", ["tiles", "tile_w", "tile_h", "image_w", "image_h", "overlap"])
|
||||
|
||||
|
||||
def split_grid(image, tile_w=512, tile_h=512, overlap=64):
|
||||
w = image.width
|
||||
h = image.height
|
||||
|
||||
non_overlap_width = tile_w - overlap
|
||||
non_overlap_height = tile_h - overlap
|
||||
|
||||
cols = math.ceil((w - overlap) / non_overlap_width)
|
||||
rows = math.ceil((h - overlap) / non_overlap_height)
|
||||
|
||||
dx = (w - tile_w) / (cols-1) if cols > 1 else 0
|
||||
dy = (h - tile_h) / (rows-1) if rows > 1 else 0
|
||||
|
||||
grid = Grid([], tile_w, tile_h, w, h, overlap)
|
||||
for row in range(rows):
|
||||
row_images = []
|
||||
|
||||
y = int(row * dy)
|
||||
|
||||
if y + tile_h >= h:
|
||||
y = h - tile_h
|
||||
|
||||
for col in range(cols):
|
||||
x = int(col * dx)
|
||||
|
||||
if x+tile_w >= w:
|
||||
x = w - tile_w
|
||||
|
||||
tile = image.crop((x, y, x + tile_w, y + tile_h))
|
||||
|
||||
row_images.append([x, tile_w, tile])
|
||||
|
||||
grid.tiles.append([y, tile_h, row_images])
|
||||
|
||||
return grid
|
||||
|
||||
|
||||
def combine_grid(grid):
|
||||
def make_mask_image(r):
|
||||
r = r * 255 / grid.overlap
|
||||
r = r.astype(np.uint8)
|
||||
return Image.fromarray(r, 'L')
|
||||
|
||||
mask_w = make_mask_image(np.arange(grid.overlap, dtype=np.float32).reshape((1, grid.overlap)).repeat(grid.tile_h, axis=0))
|
||||
mask_h = make_mask_image(np.arange(grid.overlap, dtype=np.float32).reshape((grid.overlap, 1)).repeat(grid.image_w, axis=1))
|
||||
|
||||
combined_image = Image.new("RGB", (grid.image_w, grid.image_h))
|
||||
for y, h, row in grid.tiles:
|
||||
combined_row = Image.new("RGB", (grid.image_w, h))
|
||||
for x, w, tile in row:
|
||||
if x == 0:
|
||||
combined_row.paste(tile, (0, 0))
|
||||
continue
|
||||
|
||||
combined_row.paste(tile.crop((0, 0, grid.overlap, h)), (x, 0), mask=mask_w)
|
||||
combined_row.paste(tile.crop((grid.overlap, 0, w, h)), (x + grid.overlap, 0))
|
||||
|
||||
if y == 0:
|
||||
combined_image.paste(combined_row, (0, 0))
|
||||
continue
|
||||
|
||||
combined_image.paste(combined_row.crop((0, 0, combined_row.width, grid.overlap)), (0, y), mask=mask_h)
|
||||
combined_image.paste(combined_row.crop((0, grid.overlap, combined_row.width, h)), (0, y + grid.overlap))
|
||||
|
||||
return combined_image
|
||||
|
||||
|
||||
class GridAnnotation:
|
||||
def __init__(self, text='', is_active=True):
|
||||
self.text = text
|
||||
self.is_active = is_active
|
||||
self.size = None
|
||||
|
||||
|
||||
def draw_grid_annotations(im, width, height, hor_texts, ver_texts):
|
||||
def wrap(drawing, text, font, line_length):
|
||||
lines = ['']
|
||||
for word in text.split():
|
||||
line = f'{lines[-1]} {word}'.strip()
|
||||
if drawing.textlength(line, font=font) <= line_length:
|
||||
lines[-1] = line
|
||||
else:
|
||||
lines.append(word)
|
||||
return lines
|
||||
|
||||
def draw_texts(drawing, draw_x, draw_y, lines):
|
||||
for i, line in enumerate(lines):
|
||||
drawing.multiline_text((draw_x, draw_y + line.size[1] / 2), line.text, font=fnt, fill=color_active if line.is_active else color_inactive, anchor="mm", align="center")
|
||||
|
||||
if not line.is_active:
|
||||
drawing.line((draw_x - line.size[0]//2, draw_y + line.size[1]//2, draw_x + line.size[0]//2, draw_y + line.size[1]//2), fill=color_inactive, width=4)
|
||||
|
||||
draw_y += line.size[1] + line_spacing
|
||||
|
||||
fontsize = (width + height) // 25
|
||||
line_spacing = fontsize // 2
|
||||
fnt = ImageFont.truetype(opts.font, fontsize)
|
||||
color_active = (0, 0, 0)
|
||||
color_inactive = (153, 153, 153)
|
||||
|
||||
pad_left = width * 3 // 4 if len(ver_texts) > 0 else 0
|
||||
|
||||
cols = im.width // width
|
||||
rows = im.height // height
|
||||
|
||||
assert cols == len(hor_texts), f'bad number of horizontal texts: {len(hor_texts)}; must be {cols}'
|
||||
assert rows == len(ver_texts), f'bad number of vertical texts: {len(ver_texts)}; must be {rows}'
|
||||
|
||||
calc_img = Image.new("RGB", (1, 1), "white")
|
||||
calc_d = ImageDraw.Draw(calc_img)
|
||||
|
||||
for texts, allowed_width in zip(hor_texts + ver_texts, [width] * len(hor_texts) + [pad_left] * len(ver_texts)):
|
||||
items = [] + texts
|
||||
texts.clear()
|
||||
|
||||
for line in items:
|
||||
wrapped = wrap(calc_d, line.text, fnt, allowed_width)
|
||||
texts += [GridAnnotation(x, line.is_active) for x in wrapped]
|
||||
|
||||
for line in texts:
|
||||
bbox = calc_d.multiline_textbbox((0, 0), line.text, font=fnt)
|
||||
line.size = (bbox[2] - bbox[0], bbox[3] - bbox[1])
|
||||
|
||||
hor_text_heights = [sum([line.size[1] + line_spacing for line in lines]) - line_spacing for lines in hor_texts]
|
||||
ver_text_heights = [sum([line.size[1] + line_spacing for line in lines]) - line_spacing * len(lines) for lines in ver_texts]
|
||||
|
||||
pad_top = max(hor_text_heights) + line_spacing * 2
|
||||
|
||||
result = Image.new("RGB", (im.width + pad_left, im.height + pad_top), "white")
|
||||
result.paste(im, (pad_left, pad_top))
|
||||
|
||||
d = ImageDraw.Draw(result)
|
||||
|
||||
for col in range(cols):
|
||||
x = pad_left + width * col + width / 2
|
||||
y = pad_top / 2 - hor_text_heights[col] / 2
|
||||
|
||||
draw_texts(d, x, y, hor_texts[col])
|
||||
|
||||
for row in range(rows):
|
||||
x = pad_left / 2
|
||||
y = pad_top + height * row + height / 2 - ver_text_heights[row] / 2
|
||||
|
||||
draw_texts(d, x, y, ver_texts[row])
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def draw_prompt_matrix(im, width, height, all_prompts):
|
||||
prompts = all_prompts[1:]
|
||||
boundary = math.ceil(len(prompts) / 2)
|
||||
|
||||
prompts_horiz = prompts[:boundary]
|
||||
prompts_vert = prompts[boundary:]
|
||||
|
||||
hor_texts = [[GridAnnotation(x, is_active=pos & (1 << i) != 0) for i, x in enumerate(prompts_horiz)] for pos in range(1 << len(prompts_horiz))]
|
||||
ver_texts = [[GridAnnotation(x, is_active=pos & (1 << i) != 0) for i, x in enumerate(prompts_vert)] for pos in range(1 << len(prompts_vert))]
|
||||
|
||||
return draw_grid_annotations(im, width, height, hor_texts, ver_texts)
|
||||
|
||||
|
||||
def resize_image(resize_mode, im, width, height):
|
||||
if resize_mode == 0:
|
||||
res = im.resize((width, height), resample=LANCZOS)
|
||||
elif resize_mode == 1:
|
||||
ratio = width / height
|
||||
src_ratio = im.width / im.height
|
||||
|
||||
src_w = width if ratio > src_ratio else im.width * height // im.height
|
||||
src_h = height if ratio <= src_ratio else im.height * width // im.width
|
||||
|
||||
resized = im.resize((src_w, src_h), resample=LANCZOS)
|
||||
res = Image.new("RGB", (width, height))
|
||||
res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2))
|
||||
else:
|
||||
ratio = width / height
|
||||
src_ratio = im.width / im.height
|
||||
|
||||
src_w = width if ratio < src_ratio else im.width * height // im.height
|
||||
src_h = height if ratio >= src_ratio else im.height * width // im.width
|
||||
|
||||
resized = im.resize((src_w, src_h), resample=LANCZOS)
|
||||
res = Image.new("RGB", (width, height))
|
||||
res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2))
|
||||
|
||||
if ratio < src_ratio:
|
||||
fill_height = height // 2 - src_h // 2
|
||||
res.paste(resized.resize((width, fill_height), box=(0, 0, width, 0)), box=(0, 0))
|
||||
res.paste(resized.resize((width, fill_height), box=(0, resized.height, width, resized.height)), box=(0, fill_height + src_h))
|
||||
elif ratio > src_ratio:
|
||||
fill_width = width // 2 - src_w // 2
|
||||
res.paste(resized.resize((fill_width, height), box=(0, 0, 0, height)), box=(0, 0))
|
||||
res.paste(resized.resize((fill_width, height), box=(resized.width, 0, resized.width, height)), box=(fill_width + src_w, 0))
|
||||
|
||||
return res
|
||||
|
||||
|
||||
invalid_filename_chars = '<>:"/\\|?*\n'
|
||||
|
||||
|
||||
def sanitize_filename_part(text):
|
||||
return text.replace(' ', '_').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):
|
||||
if short_filename or prompt is None or seed is None:
|
||||
file_decoration = ""
|
||||
elif opts.save_to_dirs:
|
||||
file_decoration = f"-{seed}"
|
||||
else:
|
||||
file_decoration = f"-{seed}-{sanitize_filename_part(prompt)[:128]}"
|
||||
|
||||
if extension == 'png' and opts.enable_pnginfo and info is not None:
|
||||
pnginfo = PngImagePlugin.PngInfo()
|
||||
pnginfo.add_text("parameters", 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"]
|
||||
|
||||
dirname = " ".join(words[0:opts.save_to_dirs_prompt_len])
|
||||
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])
|
||||
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}"
|
||||
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)
|
||||
|
||||
target_side_length = 4000
|
||||
oversize = image.width > target_side_length or image.height > target_side_length
|
||||
if opts.export_for_4chan and (oversize or os.stat(fullfn).st_size > 4 * 1024 * 1024):
|
||||
ratio = image.width / image.height
|
||||
|
||||
if oversize and ratio > 1:
|
||||
image = image.resize((target_side_length, image.height * target_side_length // image.width), LANCZOS)
|
||||
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)
|
||||
|
||||
if opts.save_txt and info is not None:
|
||||
with open(f"{fullfn_without_extension}.txt", "w", encoding="utf8") as file:
|
||||
file.write(info + "\n")
|
||||
|
||||
|
||||
class Upscaler:
|
||||
name = "Lanczos"
|
||||
|
||||
def do_upscale(self, img):
|
||||
return img
|
||||
|
||||
def upscale(self, img, w, h):
|
||||
for i in range(3):
|
||||
if img.width >= w and img.height >= h:
|
||||
break
|
||||
|
||||
img = self.do_upscale(img)
|
||||
|
||||
if img.width != w or img.height != h:
|
||||
img = img.resize((w, h), resample=LANCZOS)
|
||||
|
||||
return img
|
||||
|
||||
|
||||
class UpscalerNone(Upscaler):
|
||||
name = "None"
|
||||
|
||||
def upscale(self, img, w, h):
|
||||
return img
|
||||
|
||||
|
||||
modules.shared.sd_upscalers.append(UpscalerNone())
|
||||
modules.shared.sd_upscalers.append(Upscaler())
|
||||
@@ -1,145 +0,0 @@
|
||||
import math
|
||||
from PIL import Image
|
||||
|
||||
from modules.processing import Processed, StableDiffusionProcessingImg2Img, process_images
|
||||
from modules.shared import opts, state
|
||||
import modules.shared as shared
|
||||
import modules.processing as processing
|
||||
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):
|
||||
is_inpaint = mode == 1
|
||||
is_loopback = mode == 2
|
||||
is_upscale = mode == 3
|
||||
|
||||
if is_inpaint:
|
||||
image = init_img_with_mask['image']
|
||||
mask = init_img_with_mask['mask']
|
||||
else:
|
||||
image = init_img
|
||||
mask = None
|
||||
|
||||
assert 0. <= denoising_strength <= 1., 'can only work with strength in [0.0, 1.0]'
|
||||
|
||||
p = StableDiffusionProcessingImg2Img(
|
||||
sd_model=shared.sd_model,
|
||||
outpath_samples=opts.outdir_samples or opts.outdir_img2img_samples,
|
||||
outpath_grids=opts.outdir_grids or opts.outdir_img2img_grids,
|
||||
prompt=prompt,
|
||||
seed=seed,
|
||||
sampler_index=sampler_index,
|
||||
batch_size=batch_size,
|
||||
n_iter=n_iter,
|
||||
steps=steps,
|
||||
cfg_scale=cfg_scale,
|
||||
width=width,
|
||||
height=height,
|
||||
restore_faces=restore_faces,
|
||||
tiling=tiling,
|
||||
init_images=[image],
|
||||
mask=mask,
|
||||
mask_blur=mask_blur,
|
||||
inpainting_fill=inpainting_fill,
|
||||
resize_mode=resize_mode,
|
||||
denoising_strength=denoising_strength,
|
||||
inpaint_full_res=inpaint_full_res,
|
||||
inpainting_mask_invert=inpainting_mask_invert,
|
||||
extra_generation_params={"Denoising Strength": denoising_strength}
|
||||
)
|
||||
|
||||
if is_loopback:
|
||||
output_images, info = None, None
|
||||
history = []
|
||||
initial_seed = None
|
||||
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
|
||||
|
||||
upscaler = shared.sd_upscalers[upscaler_index]
|
||||
img = upscaler.upscale(init_img, init_img.width * 2, init_img.height * 2)
|
||||
|
||||
processing.torch_gc()
|
||||
|
||||
grid = images.split_grid(img, tile_w=width, tile_h=height, overlap=upscale_overlap)
|
||||
|
||||
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
|
||||
|
||||
for i in range(batch_count):
|
||||
p.init_images = work[i*p.batch_size:(i+1)*p.batch_size]
|
||||
|
||||
state.job = f"Batch {i + 1} out of {batch_count}"
|
||||
processed = process_images(p)
|
||||
|
||||
if initial_seed is None:
|
||||
initial_seed = processed.seed
|
||||
initial_info = processed.info
|
||||
|
||||
p.seed = processed.seed + 1
|
||||
work_results += processed.images
|
||||
|
||||
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
|
||||
|
||||
combined_image = images.combine_grid(grid)
|
||||
|
||||
if opts.samples_save:
|
||||
images.save_image(combined_image, p.outpath_samples, "", initial_seed, prompt, opts.grid_format, info=initial_info)
|
||||
|
||||
processed = Processed(p, [combined_image], initial_seed, initial_info)
|
||||
|
||||
else:
|
||||
|
||||
processed = modules.scripts.scripts_img2img.run(p, *args)
|
||||
|
||||
if processed is None:
|
||||
processed = process_images(p)
|
||||
|
||||
|
||||
return processed.images, processed.js(), plaintext_to_html(processed.info)
|
||||
@@ -1,73 +0,0 @@
|
||||
import torch
|
||||
|
||||
module_in_gpu = None
|
||||
cpu = torch.device("cpu")
|
||||
gpu = torch.device("cuda")
|
||||
device = gpu if torch.cuda.is_available() else cpu
|
||||
|
||||
|
||||
def setup_for_low_vram(sd_model, use_medvram):
|
||||
parents = {}
|
||||
|
||||
def send_me_to_gpu(module, _):
|
||||
"""send this module to GPU; send whatever tracked module was previous in GPU to CPU;
|
||||
we add this as forward_pre_hook to a lot of modules and this way all but one of them will
|
||||
be in CPU
|
||||
"""
|
||||
global module_in_gpu
|
||||
|
||||
module = parents.get(module, module)
|
||||
|
||||
if module_in_gpu == module:
|
||||
return
|
||||
|
||||
if module_in_gpu is not None:
|
||||
module_in_gpu.to(cpu)
|
||||
|
||||
module.to(gpu)
|
||||
module_in_gpu = module
|
||||
|
||||
# see below for register_forward_pre_hook;
|
||||
# first_stage_model does not use forward(), it uses encode/decode, so register_forward_pre_hook is
|
||||
# useless here, and we just replace those methods
|
||||
def first_stage_model_encode_wrap(self, encoder, x):
|
||||
send_me_to_gpu(self, None)
|
||||
return encoder(x)
|
||||
|
||||
def first_stage_model_decode_wrap(self, decoder, z):
|
||||
send_me_to_gpu(self, None)
|
||||
return decoder(z)
|
||||
|
||||
# remove three big modules, cond, first_stage, and unet from the model and then
|
||||
# send the model to GPU. Then put modules back. the modules will be in CPU.
|
||||
stored = sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.model
|
||||
sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.model = None, None, None
|
||||
sd_model.to(device)
|
||||
sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.model = stored
|
||||
|
||||
# register hooks for those the first two models
|
||||
sd_model.cond_stage_model.transformer.register_forward_pre_hook(send_me_to_gpu)
|
||||
sd_model.first_stage_model.register_forward_pre_hook(send_me_to_gpu)
|
||||
sd_model.first_stage_model.encode = lambda x, en=sd_model.first_stage_model.encode: first_stage_model_encode_wrap(sd_model.first_stage_model, en, x)
|
||||
sd_model.first_stage_model.decode = lambda z, de=sd_model.first_stage_model.decode: first_stage_model_decode_wrap(sd_model.first_stage_model, de, z)
|
||||
parents[sd_model.cond_stage_model.transformer] = sd_model.cond_stage_model
|
||||
|
||||
if use_medvram:
|
||||
sd_model.model.register_forward_pre_hook(send_me_to_gpu)
|
||||
else:
|
||||
diff_model = sd_model.model.diffusion_model
|
||||
|
||||
# the third remaining model is still too big for 4 GB, so we also do the same for its submodules
|
||||
# so that only one of them is in GPU at a time
|
||||
stored = diff_model.input_blocks, diff_model.middle_block, diff_model.output_blocks, diff_model.time_embed
|
||||
diff_model.input_blocks, diff_model.middle_block, diff_model.output_blocks, diff_model.time_embed = None, None, None, None
|
||||
sd_model.model.to(device)
|
||||
diff_model.input_blocks, diff_model.middle_block, diff_model.output_blocks, diff_model.time_embed = stored
|
||||
|
||||
# install hooks for bits of third model
|
||||
diff_model.time_embed.register_forward_pre_hook(send_me_to_gpu)
|
||||
for block in diff_model.input_blocks:
|
||||
block.register_forward_pre_hook(send_me_to_gpu)
|
||||
diff_model.middle_block.register_forward_pre_hook(send_me_to_gpu)
|
||||
for block in diff_model.output_blocks:
|
||||
block.register_forward_pre_hook(send_me_to_gpu)
|
||||
@@ -1,32 +0,0 @@
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
script_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
|
||||
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')]
|
||||
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
|
||||
|
||||
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'),
|
||||
]
|
||||
|
||||
paths = {}
|
||||
|
||||
for d, must_exist, what in path_dirs:
|
||||
must_exist_path = os.path.abspath(os.path.join(script_path, d, must_exist))
|
||||
if not os.path.exists(must_exist_path):
|
||||
print(f"Warning: {what} not found at path {must_exist_path}", file=sys.stderr)
|
||||
else:
|
||||
d = os.path.abspath(d)
|
||||
sys.path.append(d)
|
||||
paths[what] = d
|
||||
@@ -1,417 +0,0 @@
|
||||
import contextlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
from PIL import Image, ImageFilter, ImageOps
|
||||
import random
|
||||
|
||||
import modules.sd_hijack
|
||||
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
|
||||
|
||||
# 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()
|
||||
|
||||
|
||||
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):
|
||||
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.seed: int = seed
|
||||
self.sampler_index: int = sampler_index
|
||||
self.batch_size: int = batch_size
|
||||
self.n_iter: int = n_iter
|
||||
self.steps: int = steps
|
||||
self.cfg_scale: float = cfg_scale
|
||||
self.width: int = width
|
||||
self.height: int = height
|
||||
self.restore_faces: bool = restore_faces
|
||||
self.tiling: bool = tiling
|
||||
self.do_not_save_samples: bool = do_not_save_samples
|
||||
self.do_not_save_grid: bool = do_not_save_grid
|
||||
self.extra_generation_params: dict = extra_generation_params
|
||||
self.overlay_images = overlay_images
|
||||
self.paste_to = None
|
||||
|
||||
def init(self):
|
||||
pass
|
||||
|
||||
def sample(self, x, conditioning, unconditional_conditioning):
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class Processed:
|
||||
def __init__(self, p: StableDiffusionProcessing, images_list, seed, info):
|
||||
self.images = images_list
|
||||
self.prompt = p.prompt
|
||||
self.seed = seed
|
||||
self.info = info
|
||||
self.width = p.width
|
||||
self.height = p.height
|
||||
self.sampler = samplers[p.sampler_index].name
|
||||
self.cfg_scale = p.cfg_scale
|
||||
self.steps = p.steps
|
||||
|
||||
def js(self):
|
||||
obj = {
|
||||
"prompt": self.prompt if type(self.prompt) != list else self.prompt[0],
|
||||
"seed": int(self.seed if type(self.seed) != list else self.seed[0]),
|
||||
"width": self.width,
|
||||
"height": self.height,
|
||||
"sampler": self.sampler,
|
||||
"cfg_scale": self.cfg_scale,
|
||||
"steps": self.steps,
|
||||
}
|
||||
|
||||
return json.dumps(obj)
|
||||
|
||||
|
||||
def create_random_tensors(shape, seeds):
|
||||
xs = []
|
||||
for seed in seeds:
|
||||
torch.manual_seed(seed)
|
||||
|
||||
# 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
|
||||
# it will break everyone's seeds.
|
||||
xs.append(torch.randn(shape, device=shared.device))
|
||||
x = torch.stack(xs)
|
||||
return x
|
||||
|
||||
|
||||
def set_seed(seed):
|
||||
return int(random.randrange(4294967294)) if seed is None or seed == -1 else seed
|
||||
|
||||
|
||||
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()
|
||||
|
||||
seed = set_seed(p.seed)
|
||||
|
||||
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 = []
|
||||
|
||||
if type(prompt) == list:
|
||||
all_prompts = prompt
|
||||
else:
|
||||
all_prompts = p.batch_size * p.n_iter * [prompt]
|
||||
|
||||
if type(seed) == list:
|
||||
all_seeds = seed
|
||||
else:
|
||||
all_seeds = [int(seed + x) for x in range(len(all_prompts))]
|
||||
|
||||
def infotext(iteration=0, position_in_batch=0):
|
||||
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],
|
||||
"Face restoration": (opts.face_restoration_model if p.restore_faces else None),
|
||||
"Batch size": (None if p.batch_size < 2 else p.batch_size),
|
||||
"Batch pos": (None if p.batch_size < 2 else position_in_batch),
|
||||
}
|
||||
|
||||
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])
|
||||
|
||||
return f"{p.prompt_for_display or prompt}\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)
|
||||
|
||||
output_images = []
|
||||
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()
|
||||
|
||||
if state.job_count == -1:
|
||||
state.job_count = p.n_iter
|
||||
|
||||
for n in range(p.n_iter):
|
||||
if state.interrupted:
|
||||
break
|
||||
|
||||
prompts = all_prompts[n * p.batch_size:(n + 1) * p.batch_size]
|
||||
seeds = all_seeds[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)
|
||||
|
||||
if len(model_hijack.comments) > 0:
|
||||
comments += model_hijack.comments
|
||||
|
||||
# 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)
|
||||
|
||||
if p.n_iter > 1:
|
||||
shared.state.job = f"Batch {n+1} out of {p.n_iter}"
|
||||
|
||||
samples_ddim = p.sample(x=x, conditioning=c, unconditional_conditioning=uc)
|
||||
if state.interrupted:
|
||||
|
||||
# if we are interruped, sample returns just noise
|
||||
# use the image collected previously in sampler loop
|
||||
samples_ddim = shared.state.current_latent
|
||||
|
||||
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)
|
||||
|
||||
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()
|
||||
|
||||
x_sample = modules.face_restoration.restore_faces(x_sample)
|
||||
|
||||
image = Image.fromarray(x_sample)
|
||||
|
||||
if p.overlay_images is not None and i < len(p.overlay_images):
|
||||
overlay = p.overlay_images[i]
|
||||
|
||||
if p.paste_to is not None:
|
||||
x, y, w, h = p.paste_to
|
||||
base_image = Image.new('RGBA', (overlay.width, overlay.height))
|
||||
image = images.resize_image(1, image, w, h)
|
||||
base_image.paste(image, (x, y))
|
||||
image = base_image
|
||||
|
||||
image = image.convert('RGBA')
|
||||
image.alpha_composite(overlay)
|
||||
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))
|
||||
|
||||
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
|
||||
|
||||
grid = images.image_grid(output_images, p.batch_size)
|
||||
|
||||
if 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)
|
||||
|
||||
torch_gc()
|
||||
return Processed(p, output_images, seed, infotext())
|
||||
|
||||
|
||||
class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
|
||||
sampler = None
|
||||
|
||||
def init(self):
|
||||
self.sampler = samplers[self.sampler_index].constructor(self.sd_model)
|
||||
|
||||
def sample(self, x, conditioning, unconditional_conditioning):
|
||||
samples_ddim = self.sampler.sample(self, x, conditioning, unconditional_conditioning)
|
||||
return samples_ddim
|
||||
|
||||
|
||||
def get_crop_region(mask, pad=0):
|
||||
h, w = mask.shape
|
||||
|
||||
crop_left = 0
|
||||
for i in range(w):
|
||||
if not (mask[:, i] == 0).all():
|
||||
break
|
||||
crop_left += 1
|
||||
|
||||
crop_right = 0
|
||||
for i in reversed(range(w)):
|
||||
if not (mask[:, i] == 0).all():
|
||||
break
|
||||
crop_right += 1
|
||||
|
||||
crop_top = 0
|
||||
for i in range(h):
|
||||
if not (mask[i] == 0).all():
|
||||
break
|
||||
crop_top += 1
|
||||
|
||||
crop_bottom = 0
|
||||
for i in reversed(range(h)):
|
||||
if not (mask[i] == 0).all():
|
||||
break
|
||||
crop_bottom += 1
|
||||
|
||||
return (
|
||||
int(max(crop_left-pad, 0)),
|
||||
int(max(crop_top-pad, 0)),
|
||||
int(min(w - crop_right + pad, w)),
|
||||
int(min(h - crop_bottom + pad, h))
|
||||
)
|
||||
|
||||
|
||||
def fill(image, mask):
|
||||
image_mod = Image.new('RGBA', (image.width, image.height))
|
||||
|
||||
image_masked = Image.new('RGBa', (image.width, image.height))
|
||||
image_masked.paste(image.convert("RGBA").convert("RGBa"), mask=ImageOps.invert(mask.convert('L')))
|
||||
|
||||
image_masked = image_masked.convert('RGBa')
|
||||
|
||||
for radius, repeats in [(256, 1), (64, 1), (16, 2), (4, 4), (2, 2), (0, 1)]:
|
||||
blurred = image_masked.filter(ImageFilter.GaussianBlur(radius)).convert('RGBA')
|
||||
for _ in range(repeats):
|
||||
image_mod.alpha_composite(blurred)
|
||||
|
||||
return image_mod.convert("RGB")
|
||||
|
||||
|
||||
class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
|
||||
sampler = None
|
||||
|
||||
def __init__(self, init_images=None, resize_mode=0, denoising_strength=0.75, mask=None, mask_blur=4, inpainting_fill=0, inpaint_full_res=True, inpainting_mask_invert=0, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
self.init_images = init_images
|
||||
self.resize_mode: int = resize_mode
|
||||
self.denoising_strength: float = denoising_strength
|
||||
self.init_latent = None
|
||||
self.image_mask = mask
|
||||
#self.image_unblurred_mask = None
|
||||
self.latent_mask = None
|
||||
self.mask_for_overlay = None
|
||||
self.mask_blur = mask_blur
|
||||
self.inpainting_fill = inpainting_fill
|
||||
self.inpaint_full_res = inpaint_full_res
|
||||
self.inpainting_mask_invert = inpainting_mask_invert
|
||||
self.mask = None
|
||||
self.nmask = None
|
||||
|
||||
def init(self):
|
||||
self.sampler = samplers_for_img2img[self.sampler_index].constructor(self.sd_model)
|
||||
crop_region = None
|
||||
|
||||
if self.image_mask is not None:
|
||||
self.image_mask = self.image_mask.convert('L')
|
||||
|
||||
if self.inpainting_mask_invert:
|
||||
self.image_mask = ImageOps.invert(self.image_mask)
|
||||
|
||||
#self.image_unblurred_mask = self.image_mask
|
||||
|
||||
if self.mask_blur > 0:
|
||||
self.image_mask = self.image_mask.filter(ImageFilter.GaussianBlur(self.mask_blur))
|
||||
|
||||
if self.inpaint_full_res:
|
||||
self.mask_for_overlay = self.image_mask
|
||||
mask = self.image_mask.convert('L')
|
||||
crop_region = get_crop_region(np.array(mask), opts.upscale_at_full_resolution_padding)
|
||||
x1, y1, x2, y2 = crop_region
|
||||
|
||||
mask = mask.crop(crop_region)
|
||||
self.image_mask = images.resize_image(2, mask, self.width, self.height)
|
||||
self.paste_to = (x1, y1, x2-x1, y2-y1)
|
||||
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)
|
||||
self.mask_for_overlay = Image.fromarray(np_mask)
|
||||
|
||||
self.overlay_images = []
|
||||
|
||||
imgs = []
|
||||
for img in self.init_images:
|
||||
image = img.convert("RGB")
|
||||
|
||||
if crop_region is None:
|
||||
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')))
|
||||
|
||||
self.overlay_images.append(image_masked.convert('RGBA'))
|
||||
|
||||
if crop_region is not None:
|
||||
image = image.crop(crop_region)
|
||||
image = images.resize_image(2, image, self.width, self.height)
|
||||
|
||||
image = np.array(image).astype(np.float32) / 255.0
|
||||
image = np.moveaxis(image, 2, 0)
|
||||
|
||||
imgs.append(image)
|
||||
|
||||
if len(imgs) == 1:
|
||||
batch_images = np.expand_dims(imgs[0], axis=0).repeat(self.batch_size, axis=0)
|
||||
if self.overlay_images is not None:
|
||||
self.overlay_images = self.overlay_images * self.batch_size
|
||||
elif len(imgs) <= self.batch_size:
|
||||
self.batch_size = len(imgs)
|
||||
batch_images = np.array(imgs)
|
||||
else:
|
||||
raise RuntimeError(f"bad number of images passed: {len(imgs)}; expecting {self.batch_size} or less")
|
||||
|
||||
image = torch.from_numpy(batch_images)
|
||||
image = 2. * image - 1.
|
||||
image = image.to(shared.device)
|
||||
|
||||
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
|
||||
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 = latmask[0]
|
||||
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
|
||||
elif self.inpainting_fill == 3:
|
||||
self.init_latent = self.init_latent * self.mask
|
||||
|
||||
def sample(self, x, conditioning, unconditional_conditioning):
|
||||
samples = self.sampler.sample_img2img(self, self.init_latent, x, conditioning, unconditional_conditioning)
|
||||
|
||||
if self.mask is not None:
|
||||
samples = samples * self.nmask + self.init_latent * self.mask
|
||||
|
||||
return samples
|
||||
@@ -1,85 +0,0 @@
|
||||
import sys
|
||||
import traceback
|
||||
from collections import namedtuple
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
import modules.images
|
||||
from modules.shared import cmd_opts
|
||||
|
||||
RealesrganModelInfo = namedtuple("RealesrganModelInfo", ["name", "location", "model", "netscale"])
|
||||
|
||||
realesrgan_models = []
|
||||
have_realesrgan = False
|
||||
RealESRGANer_constructor = None
|
||||
|
||||
|
||||
class UpscalerRealESRGAN(modules.images.Upscaler):
|
||||
def __init__(self, upscaling, model_index):
|
||||
self.upscaling = upscaling
|
||||
self.model_index = model_index
|
||||
self.name = realesrgan_models[model_index].name
|
||||
|
||||
def do_upscale(self, img):
|
||||
return upscale_with_realesrgan(img, self.upscaling, self.model_index)
|
||||
|
||||
|
||||
def setup_realesrgan():
|
||||
global realesrgan_models
|
||||
global have_realesrgan
|
||||
global RealESRGANer_constructor
|
||||
|
||||
try:
|
||||
from basicsr.archs.rrdbnet_arch import RRDBNet
|
||||
from realesrgan import RealESRGANer
|
||||
from realesrgan.archs.srvgg_arch import SRVGGNetCompact
|
||||
|
||||
realesrgan_models = [
|
||||
RealesrganModelInfo(
|
||||
name="Real-ESRGAN 4x plus",
|
||||
location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
|
||||
netscale=4, model=lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
|
||||
),
|
||||
RealesrganModelInfo(
|
||||
name="Real-ESRGAN 4x plus anime 6B",
|
||||
location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth",
|
||||
netscale=4, model=lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
|
||||
),
|
||||
RealesrganModelInfo(
|
||||
name="Real-ESRGAN 2x plus",
|
||||
location="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
|
||||
netscale=2, model=lambda: RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
|
||||
),
|
||||
]
|
||||
have_realesrgan = True
|
||||
RealESRGANer_constructor = RealESRGANer
|
||||
|
||||
for i, model in enumerate(realesrgan_models):
|
||||
modules.shared.sd_upscalers.append(UpscalerRealESRGAN(model.netscale, i))
|
||||
|
||||
except Exception:
|
||||
print("Error importing Real-ESRGAN:", file=sys.stderr)
|
||||
print(traceback.format_exc(), file=sys.stderr)
|
||||
|
||||
realesrgan_models = [RealesrganModelInfo('None', '', 0, None)]
|
||||
have_realesrgan = False
|
||||
|
||||
|
||||
def upscale_with_realesrgan(image, RealESRGAN_upscaling, RealESRGAN_model_index):
|
||||
if not have_realesrgan or RealESRGANer_constructor is None:
|
||||
return image
|
||||
|
||||
info = realesrgan_models[RealESRGAN_model_index]
|
||||
|
||||
model = info.model()
|
||||
upsampler = RealESRGANer_constructor(
|
||||
scale=info.netscale,
|
||||
model_path=info.location,
|
||||
model=model,
|
||||
half=not cmd_opts.no_half
|
||||
)
|
||||
|
||||
upsampled = upsampler.enhance(np.array(image), outscale=RealESRGAN_upscaling)[0]
|
||||
|
||||
image = Image.fromarray(upsampled)
|
||||
return image
|
||||
@@ -1,144 +0,0 @@
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
|
||||
import modules.ui as ui
|
||||
import gradio as gr
|
||||
|
||||
from modules.processing import StableDiffusionProcessing
|
||||
|
||||
class Script:
|
||||
filename = None
|
||||
args_from = None
|
||||
args_to = None
|
||||
|
||||
def title(self):
|
||||
raise NotImplementedError()
|
||||
|
||||
def ui(self, is_img2img):
|
||||
pass
|
||||
|
||||
def show(self, is_img2img):
|
||||
return True
|
||||
|
||||
def run(self, *args):
|
||||
raise NotImplementedError()
|
||||
|
||||
def describe(self):
|
||||
return ""
|
||||
|
||||
|
||||
scripts_data = []
|
||||
|
||||
|
||||
def load_scripts(basedir):
|
||||
if not os.path.exists(basedir):
|
||||
return
|
||||
|
||||
for filename in os.listdir(basedir):
|
||||
path = os.path.join(basedir, filename)
|
||||
|
||||
if not os.path.isfile(path):
|
||||
continue
|
||||
|
||||
with open(path, "r", encoding="utf8") as file:
|
||||
text = file.read()
|
||||
|
||||
try:
|
||||
from types import ModuleType
|
||||
compiled = compile(text, path, 'exec')
|
||||
module = ModuleType(filename)
|
||||
exec(compiled, module.__dict__)
|
||||
|
||||
for key, script_class in module.__dict__.items():
|
||||
if type(script_class) == type and issubclass(script_class, Script):
|
||||
scripts_data.append((script_class, path))
|
||||
|
||||
except Exception:
|
||||
print(f"Error loading script: {filename}", file=sys.stderr)
|
||||
print(traceback.format_exc(), file=sys.stderr)
|
||||
|
||||
|
||||
def wrap_call(func, filename, funcname, *args, default=None, **kwargs):
|
||||
try:
|
||||
res = func(*args, **kwargs)
|
||||
return res
|
||||
except Exception:
|
||||
print(f"Error calling: {filename}/{funcname}", file=sys.stderr)
|
||||
print(traceback.format_exc(), file=sys.stderr)
|
||||
|
||||
return default
|
||||
|
||||
|
||||
class ScriptRunner:
|
||||
def __init__(self):
|
||||
self.scripts = []
|
||||
|
||||
def setup_ui(self, is_img2img):
|
||||
for script_class, path in scripts_data:
|
||||
script = script_class()
|
||||
script.filename = path
|
||||
|
||||
if not script.show(is_img2img):
|
||||
continue
|
||||
|
||||
self.scripts.append(script)
|
||||
|
||||
titles = [wrap_call(script.title, script.filename, "title") or f"{script.filename} [error]" for script in self.scripts]
|
||||
|
||||
dropdown = gr.Dropdown(label="Script", choices=["None"] + titles, value="None", type="index")
|
||||
inputs = [dropdown]
|
||||
|
||||
for script in self.scripts:
|
||||
script.args_from = len(inputs)
|
||||
|
||||
controls = wrap_call(script.ui, script.filename, "ui", is_img2img)
|
||||
|
||||
if controls is None:
|
||||
continue
|
||||
|
||||
for control in controls:
|
||||
control.visible = False
|
||||
|
||||
inputs += controls
|
||||
script.args_to = len(inputs)
|
||||
|
||||
def select_script(script_index):
|
||||
if 0 < script_index <= len(self.scripts):
|
||||
script = self.scripts[script_index-1]
|
||||
args_from = script.args_from
|
||||
args_to = script.args_to
|
||||
else:
|
||||
args_from = 0
|
||||
args_to = 0
|
||||
|
||||
return [ui.gr_show(True if i == 0 else args_from <= i < args_to) for i in range(len(inputs))]
|
||||
|
||||
dropdown.change(
|
||||
fn=select_script,
|
||||
inputs=[dropdown],
|
||||
outputs=inputs
|
||||
)
|
||||
|
||||
return inputs
|
||||
|
||||
|
||||
def run(self, p: StableDiffusionProcessing, *args):
|
||||
script_index = args[0]
|
||||
|
||||
if script_index == 0:
|
||||
return None
|
||||
|
||||
script = self.scripts[script_index-1]
|
||||
|
||||
if script is None:
|
||||
return None
|
||||
|
||||
script_args = args[script.args_from:script.args_to]
|
||||
processed = script.run(p, *script_args)
|
||||
|
||||
return processed
|
||||
|
||||
|
||||
scripts_txt2img = ScriptRunner()
|
||||
scripts_img2img = ScriptRunner()
|
||||
@@ -1,275 +0,0 @@
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
import torch
|
||||
import numpy as np
|
||||
from torch import einsum
|
||||
|
||||
from modules.shared import opts, device, cmd_opts
|
||||
|
||||
from ldm.util import default
|
||||
from einops import rearrange
|
||||
import ldm.modules.attention
|
||||
|
||||
|
||||
# see https://github.com/basujindal/stable-diffusion/pull/117 for discussion
|
||||
def split_cross_attention_forward(self, x, context=None, mask=None):
|
||||
h = self.heads
|
||||
|
||||
q = self.to_q(x)
|
||||
context = default(context, x)
|
||||
k = self.to_k(context)
|
||||
v = 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, k, v))
|
||||
|
||||
r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device)
|
||||
for i in range(0, q.shape[0], 2):
|
||||
end = i + 2
|
||||
s1 = einsum('b i d, b j d -> b i j', q[i:end], k[i:end])
|
||||
s1 *= self.scale
|
||||
|
||||
s2 = s1.softmax(dim=-1)
|
||||
del s1
|
||||
|
||||
r1[i:end] = einsum('b i j, b j d -> b i d', s2, v[i:end])
|
||||
del s2
|
||||
|
||||
r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h)
|
||||
del r1
|
||||
|
||||
return self.to_out(r2)
|
||||
|
||||
|
||||
class StableDiffusionModelHijack:
|
||||
ids_lookup = {}
|
||||
word_embeddings = {}
|
||||
word_embeddings_checksums = {}
|
||||
fixes = None
|
||||
comments = []
|
||||
dir_mtime = None
|
||||
layers = None
|
||||
circular_enabled = False
|
||||
|
||||
def load_textual_inversion_embeddings(self, dirname, model):
|
||||
mt = os.path.getmtime(dirname)
|
||||
if self.dir_mtime is not None and mt <= self.dir_mtime:
|
||||
return
|
||||
|
||||
self.dir_mtime = mt
|
||||
self.ids_lookup.clear()
|
||||
self.word_embeddings.clear()
|
||||
|
||||
tokenizer = model.cond_stage_model.tokenizer
|
||||
|
||||
def const_hash(a):
|
||||
r = 0
|
||||
for v in a:
|
||||
r = (r * 281 ^ int(v) * 997) & 0xFFFFFFFF
|
||||
return r
|
||||
|
||||
def process_file(path, filename):
|
||||
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]
|
||||
self.word_embeddings[name] = emb.detach()
|
||||
self.word_embeddings_checksums[name] = f'{const_hash(emb.reshape(-1))&0xffff:04x}'
|
||||
|
||||
ids = tokenizer([name], add_special_tokens=False)['input_ids'][0]
|
||||
|
||||
first_id = ids[0]
|
||||
if first_id not in self.ids_lookup:
|
||||
self.ids_lookup[first_id] = []
|
||||
self.ids_lookup[first_id].append((ids, name))
|
||||
|
||||
for fn in os.listdir(dirname):
|
||||
try:
|
||||
process_file(os.path.join(dirname, fn), fn)
|
||||
except Exception:
|
||||
print(f"Error loading emedding {fn}:", file=sys.stderr)
|
||||
print(traceback.format_exc(), file=sys.stderr)
|
||||
continue
|
||||
|
||||
print(f"Loaded a total of {len(self.word_embeddings)} text inversion embeddings.")
|
||||
|
||||
def hijack(self, m):
|
||||
model_embeddings = m.cond_stage_model.transformer.text_model.embeddings
|
||||
|
||||
model_embeddings.token_embedding = EmbeddingsWithFixes(model_embeddings.token_embedding, self)
|
||||
m.cond_stage_model = FrozenCLIPEmbedderWithCustomWords(m.cond_stage_model, self)
|
||||
|
||||
if cmd_opts.opt_split_attention:
|
||||
ldm.modules.attention.CrossAttention.forward = split_cross_attention_forward
|
||||
|
||||
def flatten(el):
|
||||
flattened = [flatten(children) for children in el.children()]
|
||||
res = [el]
|
||||
for c in flattened:
|
||||
res += c
|
||||
return res
|
||||
|
||||
self.layers = flatten(m)
|
||||
|
||||
def apply_circular(self, enable):
|
||||
if self.circular_enabled == enable:
|
||||
return
|
||||
|
||||
self.circular_enabled = enable
|
||||
|
||||
for layer in [layer for layer in self.layers if type(layer) == torch.nn.Conv2d]:
|
||||
layer.padding_mode = 'circular' if enable else 'zeros'
|
||||
|
||||
|
||||
class FrozenCLIPEmbedderWithCustomWords(torch.nn.Module):
|
||||
def __init__(self, wrapped, hijack):
|
||||
super().__init__()
|
||||
self.wrapped = wrapped
|
||||
self.hijack = hijack
|
||||
self.tokenizer = wrapped.tokenizer
|
||||
self.max_length = wrapped.max_length
|
||||
self.token_mults = {}
|
||||
|
||||
tokens_with_parens = [(k, v) for k, v in self.tokenizer.get_vocab().items() if '(' in k or ')' in k or '[' in k or ']' in k]
|
||||
for text, ident in tokens_with_parens:
|
||||
mult = 1.0
|
||||
for c in text:
|
||||
if c == '[':
|
||||
mult /= 1.1
|
||||
if c == ']':
|
||||
mult *= 1.1
|
||||
if c == '(':
|
||||
mult *= 1.1
|
||||
if c == ')':
|
||||
mult /= 1.1
|
||||
|
||||
if mult != 1.0:
|
||||
self.token_mults[ident] = mult
|
||||
|
||||
def forward(self, text):
|
||||
self.hijack.fixes = []
|
||||
self.hijack.comments = []
|
||||
remade_batch_tokens = []
|
||||
id_start = self.wrapped.tokenizer.bos_token_id
|
||||
id_end = self.wrapped.tokenizer.eos_token_id
|
||||
maxlen = self.wrapped.max_length - 2
|
||||
used_custom_terms = []
|
||||
|
||||
cache = {}
|
||||
batch_tokens = self.wrapped.tokenizer(text, truncation=False, add_special_tokens=False)["input_ids"]
|
||||
batch_multipliers = []
|
||||
for tokens in batch_tokens:
|
||||
tuple_tokens = tuple(tokens)
|
||||
|
||||
if tuple_tokens in cache:
|
||||
remade_tokens, fixes, multipliers = cache[tuple_tokens]
|
||||
else:
|
||||
fixes = []
|
||||
remade_tokens = []
|
||||
multipliers = []
|
||||
mult = 1.0
|
||||
|
||||
i = 0
|
||||
while i < len(tokens):
|
||||
token = tokens[i]
|
||||
|
||||
possible_matches = self.hijack.ids_lookup.get(token, None)
|
||||
|
||||
mult_change = self.token_mults.get(token) if opts.enable_emphasis else None
|
||||
if mult_change is not None:
|
||||
mult *= mult_change
|
||||
elif possible_matches is None:
|
||||
remade_tokens.append(token)
|
||||
multipliers.append(mult)
|
||||
else:
|
||||
found = False
|
||||
for ids, word in possible_matches:
|
||||
if tokens[i:i+len(ids)] == ids:
|
||||
emb_len = int(self.hijack.word_embeddings[word].shape[0])
|
||||
fixes.append((len(remade_tokens), word))
|
||||
remade_tokens += [0] * emb_len
|
||||
multipliers += [mult] * emb_len
|
||||
i += len(ids) - 1
|
||||
found = True
|
||||
used_custom_terms.append((word, self.hijack.word_embeddings_checksums[word]))
|
||||
break
|
||||
|
||||
if not found:
|
||||
remade_tokens.append(token)
|
||||
multipliers.append(mult)
|
||||
|
||||
i += 1
|
||||
|
||||
if len(remade_tokens) > maxlen - 2:
|
||||
vocab = {v: k for k, v in self.wrapped.tokenizer.get_vocab().items()}
|
||||
ovf = remade_tokens[maxlen - 2:]
|
||||
overflowing_words = [vocab.get(int(x), "") for x in ovf]
|
||||
overflowing_text = self.wrapped.tokenizer.convert_tokens_to_string(''.join(overflowing_words))
|
||||
|
||||
self.hijack.comments.append(f"Warning: too many input tokens; some ({len(overflowing_words)}) have been truncated:\n{overflowing_text}\n")
|
||||
|
||||
remade_tokens = remade_tokens + [id_end] * (maxlen - 2 - len(remade_tokens))
|
||||
remade_tokens = [id_start] + remade_tokens[0:maxlen-2] + [id_end]
|
||||
cache[tuple_tokens] = (remade_tokens, fixes, multipliers)
|
||||
|
||||
multipliers = multipliers + [1.0] * (maxlen - 2 - len(multipliers))
|
||||
multipliers = [1.0] + multipliers[0:maxlen - 2] + [1.0]
|
||||
|
||||
remade_batch_tokens.append(remade_tokens)
|
||||
self.hijack.fixes.append(fixes)
|
||||
batch_multipliers.append(multipliers)
|
||||
|
||||
if len(used_custom_terms) > 0:
|
||||
self.hijack.comments.append("Used custom terms: " + ", ".join([f'{word} [{checksum}]' for word, checksum in used_custom_terms]))
|
||||
|
||||
tokens = torch.asarray(remade_batch_tokens).to(device)
|
||||
outputs = self.wrapped.transformer(input_ids=tokens)
|
||||
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)
|
||||
original_mean = z.mean()
|
||||
z *= batch_multipliers.reshape(batch_multipliers.shape + (1,)).expand(z.shape)
|
||||
new_mean = z.mean()
|
||||
z *= original_mean / new_mean
|
||||
|
||||
return z
|
||||
|
||||
|
||||
class EmbeddingsWithFixes(torch.nn.Module):
|
||||
def __init__(self, wrapped, embeddings):
|
||||
super().__init__()
|
||||
self.wrapped = wrapped
|
||||
self.embeddings = embeddings
|
||||
|
||||
def forward(self, input_ids):
|
||||
batch_fixes = self.embeddings.fixes
|
||||
self.embeddings.fixes = None
|
||||
|
||||
inputs_embeds = self.wrapped(input_ids)
|
||||
|
||||
if batch_fixes is not None:
|
||||
for fixes, tensor in zip(batch_fixes, inputs_embeds):
|
||||
for offset, word in fixes:
|
||||
emb = self.embeddings.word_embeddings[word]
|
||||
emb_len = min(tensor.shape[0]-offset, emb.shape[0])
|
||||
tensor[offset:offset+emb_len] = self.embeddings.word_embeddings[word][0:emb_len]
|
||||
|
||||
return inputs_embeds
|
||||
|
||||
|
||||
def add_circular_option_to_conv_2d():
|
||||
conv2d_constructor = torch.nn.Conv2d.__init__
|
||||
|
||||
def conv2d_constructor_circular(self, *args, **kwargs):
|
||||
return conv2d_constructor(self, *args, padding_mode='circular', **kwargs)
|
||||
|
||||
torch.nn.Conv2d.__init__ = conv2d_constructor_circular
|
||||
|
||||
|
||||
model_hijack = StableDiffusionModelHijack()
|
||||
@@ -1,192 +0,0 @@
|
||||
from collections import namedtuple
|
||||
import numpy as np
|
||||
import torch
|
||||
import tqdm
|
||||
from PIL import Image
|
||||
|
||||
import k_diffusion.sampling
|
||||
import ldm.models.diffusion.ddim
|
||||
import ldm.models.diffusion.plms
|
||||
|
||||
from modules.shared import opts, cmd_opts, state
|
||||
import modules.shared as shared
|
||||
|
||||
|
||||
SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases'])
|
||||
|
||||
samplers_k_diffusion = [
|
||||
('Euler a', 'sample_euler_ancestral', ['k_euler_a']),
|
||||
('Euler', 'sample_euler', ['k_euler']),
|
||||
('LMS', 'sample_lms', ['k_lms']),
|
||||
('Heun', 'sample_heun', ['k_heun']),
|
||||
('DPM2', 'sample_dpm_2', ['k_dpm_2']),
|
||||
('DPM2 a', 'sample_dpm_2_ancestral', ['k_dpm_2_a']),
|
||||
]
|
||||
|
||||
samplers_data_k_diffusion = [
|
||||
SamplerData(label, lambda model, funcname=funcname: KDiffusionSampler(funcname, model), aliases)
|
||||
for label, funcname, aliases in samplers_k_diffusion
|
||||
if hasattr(k_diffusion.sampling, funcname)
|
||||
]
|
||||
|
||||
samplers = [
|
||||
*samplers_data_k_diffusion,
|
||||
SamplerData('DDIM', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.ddim.DDIMSampler, model), []),
|
||||
SamplerData('PLMS', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.plms.PLMSSampler, model), []),
|
||||
]
|
||||
samplers_for_img2img = [x for x in samplers if x.name != 'PLMS']
|
||||
|
||||
|
||||
def sample_to_image(samples):
|
||||
x_sample = shared.sd_model.decode_first_stage(samples[0:1].type(shared.sd_model.dtype))[0]
|
||||
x_sample = torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0)
|
||||
x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)
|
||||
x_sample = x_sample.astype(np.uint8)
|
||||
return Image.fromarray(x_sample)
|
||||
|
||||
|
||||
def store_latent(decoded):
|
||||
state.current_latent = decoded
|
||||
|
||||
if opts.show_progress_every_n_steps > 0 and shared.state.sampling_step % opts.show_progress_every_n_steps == 0:
|
||||
if not shared.parallel_processing_allowed:
|
||||
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):
|
||||
if state.interrupted:
|
||||
break
|
||||
|
||||
yield x
|
||||
|
||||
state.sampling_step += 1
|
||||
|
||||
|
||||
ldm.models.diffusion.ddim.tqdm = lambda *args, desc=None, **kwargs: extended_tdqm(*args, desc=desc, **kwargs)
|
||||
ldm.models.diffusion.plms.tqdm = lambda *args, desc=None, **kwargs: extended_tdqm(*args, desc=desc, **kwargs)
|
||||
|
||||
|
||||
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.mask = None
|
||||
self.nmask = None
|
||||
self.init_latent = None
|
||||
|
||||
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
|
||||
try:
|
||||
self.sampler.make_schedule(ddim_num_steps=p.steps, verbose=False)
|
||||
except Exception:
|
||||
self.sampler.make_schedule(ddim_num_steps=p.steps+1, verbose=False)
|
||||
|
||||
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.mask = p.mask
|
||||
self.nmask = p.nmask
|
||||
self.init_latent = p.init_latent
|
||||
|
||||
samples = self.sampler.decode(x1, conditioning, t_enc, unconditional_guidance_scale=p.cfg_scale, unconditional_conditioning=unconditional_conditioning)
|
||||
|
||||
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)
|
||||
return samples_ddim
|
||||
|
||||
|
||||
class CFGDenoiser(torch.nn.Module):
|
||||
def __init__(self, model):
|
||||
super().__init__()
|
||||
self.inner_model = model
|
||||
self.mask = None
|
||||
self.nmask = None
|
||||
self.init_latent = None
|
||||
|
||||
def forward(self, x, sigma, uncond, cond, cond_scale):
|
||||
if shared.batch_cond_uncond:
|
||||
x_in = torch.cat([x] * 2)
|
||||
sigma_in = torch.cat([sigma] * 2)
|
||||
cond_in = torch.cat([uncond, cond])
|
||||
uncond, cond = self.inner_model(x_in, sigma_in, cond=cond_in).chunk(2)
|
||||
denoised = uncond + (cond - uncond) * cond_scale
|
||||
else:
|
||||
uncond = self.inner_model(x, sigma, cond=uncond)
|
||||
cond = self.inner_model(x, sigma, cond=cond)
|
||||
denoised = uncond + (cond - uncond) * cond_scale
|
||||
|
||||
if self.mask is not None:
|
||||
denoised = self.init_latent * self.mask + self.nmask * denoised
|
||||
|
||||
return denoised
|
||||
|
||||
|
||||
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):
|
||||
if state.interrupted:
|
||||
break
|
||||
|
||||
yield x
|
||||
|
||||
state.sampling_step += 1
|
||||
|
||||
|
||||
class KDiffusionSampler:
|
||||
def __init__(self, funcname, sd_model):
|
||||
self.model_wrap = k_diffusion.external.CompVisDenoiser(sd_model)
|
||||
self.funcname = funcname
|
||||
self.func = getattr(k_diffusion.sampling, self.funcname)
|
||||
self.model_wrap_cfg = CFGDenoiser(self.model_wrap)
|
||||
|
||||
def callback_state(self, d):
|
||||
store_latent(d["denoised"])
|
||||
|
||||
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
|
||||
|
||||
sigma_sched = sigmas[p.steps - t_enc - 1:]
|
||||
|
||||
self.model_wrap_cfg.mask = p.mask
|
||||
self.model_wrap_cfg.nmask = p.nmask
|
||||
self.model_wrap_cfg.init_latent = p.init_latent
|
||||
|
||||
if hasattr(k_diffusion.sampling, 'trange'):
|
||||
k_diffusion.sampling.trange = lambda *args, **kwargs: extended_trange(*args, **kwargs)
|
||||
|
||||
return self.func(self.model_wrap_cfg, xi, sigma_sched, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': p.cfg_scale}, disable=False, callback=self.callback_state)
|
||||
|
||||
def sample(self, p, x, conditioning, unconditional_conditioning):
|
||||
sigmas = self.model_wrap.get_sigmas(p.steps)
|
||||
x = x * sigmas[0]
|
||||
|
||||
if hasattr(k_diffusion.sampling, 'trange'):
|
||||
k_diffusion.sampling.trange = lambda *args, **kwargs: extended_trange(*args, **kwargs)
|
||||
|
||||
samples_ddim = self.func(self.model_wrap_cfg, x, sigmas, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': p.cfg_scale}, disable=False, callback=self.callback_state)
|
||||
return samples_ddim
|
||||
|
||||
@@ -1,164 +0,0 @@
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
|
||||
import gradio as gr
|
||||
import torch
|
||||
|
||||
import modules.artists
|
||||
from modules.paths import script_path, sd_path
|
||||
import modules.codeformer_model
|
||||
|
||||
config_filename = "config.json"
|
||||
|
||||
sd_model_file = os.path.join(script_path, 'model.ckpt')
|
||||
if not os.path.exists(sd_model_file):
|
||||
sd_model_file = "models/ldm/stable-diffusion-v1/model.ckpt"
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--config", type=str, default=os.path.join(sd_path, "configs/stable-diffusion/v1-inference.yaml"), help="path to config which constructs model",)
|
||||
parser.add_argument("--ckpt", type=str, default=os.path.join(sd_path, sd_model_file), help="path to checkpoint of model",)
|
||||
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("--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("--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("--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("--listen", action='store_true', help="launch gradio with 0.0.0.0 as server name, allowing to respond to network requests")
|
||||
cmd_opts = parser.parse_args()
|
||||
|
||||
cpu = torch.device("cpu")
|
||||
gpu = torch.device("cuda")
|
||||
device = gpu if torch.cuda.is_available() else cpu
|
||||
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
|
||||
|
||||
|
||||
class State:
|
||||
interrupted = False
|
||||
job = ""
|
||||
job_no = 0
|
||||
job_count = 0
|
||||
sampling_step = 0
|
||||
sampling_steps = 0
|
||||
current_latent = None
|
||||
current_image = None
|
||||
current_image_sampling_step = 0
|
||||
|
||||
|
||||
def interrupt(self):
|
||||
self.interrupted = True
|
||||
|
||||
def nextjob(self):
|
||||
self.job_no += 1
|
||||
self.sampling_step = 0
|
||||
self.current_image_sampling_step = 0
|
||||
|
||||
|
||||
state = State()
|
||||
|
||||
artist_db = modules.artists.ArtistsDatabase(os.path.join(script_path, 'artists.csv'))
|
||||
|
||||
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):
|
||||
self.default = default
|
||||
self.label = label
|
||||
self.component = component
|
||||
self.component_args = component_args
|
||||
|
||||
data = 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"),
|
||||
"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"),
|
||||
"grid_only_if_multiple": OptionInfo(True, "Do not save grids consisting of one picture"),
|
||||
"n_rows": OptionInfo(-1, "Grid row count; use -1 for autodetect and 0 for it to be same as batch size", gr.Slider, {"minimum": -1, "maximum": 16, "step": 1}),
|
||||
"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"),
|
||||
"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}),
|
||||
"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}),
|
||||
"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}),
|
||||
}
|
||||
|
||||
def __init__(self):
|
||||
self.data = {k: v.default for k, v in self.data_labels.items()}
|
||||
|
||||
def __setattr__(self, key, value):
|
||||
if self.data is not None:
|
||||
if key in self.data:
|
||||
self.data[key] = value
|
||||
|
||||
return super(Options, self).__setattr__(key, value)
|
||||
|
||||
def __getattr__(self, item):
|
||||
if self.data is not None:
|
||||
if item in self.data:
|
||||
return self.data[item]
|
||||
|
||||
if item in self.data_labels:
|
||||
return self.data_labels[item].default
|
||||
|
||||
return super(Options, self).__getattribute__(item)
|
||||
|
||||
def save(self, filename):
|
||||
with open(filename, "w", encoding="utf8") as file:
|
||||
json.dump(self.data, file)
|
||||
|
||||
def load(self, filename):
|
||||
with open(filename, "r", encoding="utf8") as file:
|
||||
self.data = json.load(file)
|
||||
|
||||
|
||||
opts = Options()
|
||||
if os.path.exists(config_filename):
|
||||
opts.load(config_filename)
|
||||
|
||||
sd_upscalers = []
|
||||
|
||||
sd_model = None
|
||||
|
||||
|
||||
@@ -1,36 +0,0 @@
|
||||
import modules.scripts
|
||||
from modules.processing import StableDiffusionProcessing, Processed, StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, process_images
|
||||
from modules.shared import opts, cmd_opts
|
||||
import modules.shared as shared
|
||||
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):
|
||||
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,
|
||||
negative_prompt=negative_prompt,
|
||||
seed=seed,
|
||||
sampler_index=sampler_index,
|
||||
batch_size=batch_size,
|
||||
n_iter=n_iter,
|
||||
steps=steps,
|
||||
cfg_scale=cfg_scale,
|
||||
width=width,
|
||||
height=height,
|
||||
restore_faces=restore_faces,
|
||||
tiling=tiling,
|
||||
)
|
||||
|
||||
processed = modules.scripts.scripts_txt2img.run(p, *args)
|
||||
|
||||
if processed is not None:
|
||||
pass
|
||||
else:
|
||||
processed = process_images(p)
|
||||
|
||||
return processed.images, processed.js(), plaintext_to_html(processed.info)
|
||||
|
||||
-709
@@ -1,709 +0,0 @@
|
||||
import base64
|
||||
import html
|
||||
import io
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
import random
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
import gradio as gr
|
||||
import gradio.utils
|
||||
import gradio.routes
|
||||
|
||||
from modules.paths import script_path
|
||||
from modules.shared import opts, cmd_opts
|
||||
import modules.shared as shared
|
||||
from modules.sd_samplers import samplers, samplers_for_img2img
|
||||
import modules.realesrgan_model as realesrgan
|
||||
import modules.scripts
|
||||
import modules.gfpgan_model
|
||||
import modules.codeformer_model
|
||||
|
||||
# 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()
|
||||
mimetypes.add_type('application/javascript', '.js')
|
||||
|
||||
|
||||
if not cmd_opts.share and not cmd_opts.listen:
|
||||
# fix gradio phoning home
|
||||
gradio.utils.version_check = lambda: None
|
||||
gradio.utils.get_local_ip_address = lambda: '127.0.0.1'
|
||||
|
||||
|
||||
def gr_show(visible=True):
|
||||
return {"visible": visible, "__type__": "update"}
|
||||
|
||||
|
||||
sample_img2img = "assets/stable-samples/img2img/sketch-mountains-input.jpg"
|
||||
sample_img2img = sample_img2img if os.path.exists(sample_img2img) else None
|
||||
|
||||
css_hide_progressbar = """
|
||||
.wrap .m-12 svg { display:none!important; }
|
||||
.wrap .m-12::before { content:"Loading..." }
|
||||
.progress-bar { display:none!important; }
|
||||
.meta-text { display:none!important; }
|
||||
"""
|
||||
|
||||
def plaintext_to_html(text):
|
||||
text = "".join([f"<p>{html.escape(x)}</p>\n" for x in text.split('\n')])
|
||||
return text
|
||||
|
||||
|
||||
def image_from_url_text(filedata):
|
||||
if type(filedata) == list:
|
||||
if len(filedata) == 0:
|
||||
return None
|
||||
|
||||
filedata = filedata[0]
|
||||
|
||||
if filedata.startswith("data:image/png;base64,"):
|
||||
filedata = filedata[len("data:image/png;base64,"):]
|
||||
|
||||
filedata = base64.decodebytes(filedata.encode('utf-8'))
|
||||
image = Image.open(io.BytesIO(filedata))
|
||||
return image
|
||||
|
||||
|
||||
def send_gradio_gallery_to_image(x):
|
||||
if len(x) == 0:
|
||||
return None
|
||||
|
||||
return image_from_url_text(x[0])
|
||||
|
||||
|
||||
def save_files(js_data, images):
|
||||
import csv
|
||||
|
||||
os.makedirs(opts.outdir_save, exist_ok=True)
|
||||
|
||||
filenames = []
|
||||
|
||||
data = json.loads(js_data)
|
||||
|
||||
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"])
|
||||
|
||||
filename_base = str(int(time.time() * 1000))
|
||||
for i, filedata in enumerate(images):
|
||||
filename = filename_base + ("" if len(images) == 1 else "-" + str(i + 1)) + ".png"
|
||||
filepath = os.path.join(opts.outdir_save, filename)
|
||||
|
||||
if filedata.startswith("data:image/png;base64,"):
|
||||
filedata = filedata[len("data:image/png;base64,"):]
|
||||
|
||||
with open(filepath, "wb") as imgfile:
|
||||
imgfile.write(base64.decodebytes(filedata.encode('utf-8')))
|
||||
|
||||
filenames.append(filename)
|
||||
|
||||
writer.writerow([data["prompt"], data["seed"], data["width"], data["height"], data["sampler"], data["cfg_scale"], data["steps"], filenames[0]])
|
||||
|
||||
return '', '', plaintext_to_html(f"Saved: {filenames[0]}")
|
||||
|
||||
|
||||
def wrap_gradio_call(func):
|
||||
def f(*args, **kwargs):
|
||||
t = time.perf_counter()
|
||||
|
||||
try:
|
||||
res = list(func(*args, **kwargs))
|
||||
except Exception as e:
|
||||
print("Error completing request", file=sys.stderr)
|
||||
print("Arguments:", args, kwargs, file=sys.stderr)
|
||||
print(traceback.format_exc(), file=sys.stderr)
|
||||
|
||||
shared.state.job = ""
|
||||
shared.state.job_count = 0
|
||||
|
||||
res = [None, '', f"<div class='error'>{plaintext_to_html(type(e).__name__+': '+str(e))}</div>"]
|
||||
|
||||
elapsed = time.perf_counter() - t
|
||||
|
||||
# last item is always HTML
|
||||
res[-1] = res[-1] + f"<p class='performance'>Time taken: {elapsed:.2f}s</p>"
|
||||
|
||||
shared.state.interrupted = False
|
||||
|
||||
return tuple(res)
|
||||
|
||||
return f
|
||||
|
||||
|
||||
def check_progress_call():
|
||||
|
||||
if shared.state.job_count == 0:
|
||||
return "", gr_show(False), gr_show(False)
|
||||
|
||||
progress = 0
|
||||
|
||||
if shared.state.job_count > 0:
|
||||
progress += shared.state.job_no / shared.state.job_count
|
||||
if shared.state.sampling_steps > 0:
|
||||
progress += 1 / shared.state.job_count * shared.state.sampling_step / shared.state.sampling_steps
|
||||
|
||||
progress = min(progress, 1)
|
||||
|
||||
progressbar = ""
|
||||
if opts.show_progressbar:
|
||||
progressbar = f"""<div class='progressDiv'><div class='progress' style="width:{progress * 100}%">{str(int(progress*100))+"%" if progress > 0.01 else ""}</div></div>"""
|
||||
|
||||
image = gr_show(False)
|
||||
preview_visibility = gr_show(False)
|
||||
|
||||
if opts.show_progress_every_n_steps > 0:
|
||||
if shared.parallel_processing_allowed:
|
||||
|
||||
if shared.state.sampling_step - shared.state.current_image_sampling_step >= opts.show_progress_every_n_steps and shared.state.current_latent is not None:
|
||||
shared.state.current_image = modules.sd_samplers.sample_to_image(shared.state.current_latent)
|
||||
shared.state.current_image_sampling_step = shared.state.sampling_step
|
||||
|
||||
image = shared.state.current_image
|
||||
|
||||
if image is None or progress >= 1:
|
||||
image = gr.update(value=None)
|
||||
else:
|
||||
preview_visibility = gr_show(True)
|
||||
|
||||
return f"<span style='display: none'>{time.time()}</span><p>{progressbar}</p>", preview_visibility, image
|
||||
|
||||
|
||||
def roll_artist(prompt):
|
||||
allowed_cats = set([x for x in shared.artist_db.categories() if len(opts.random_artist_categories)==0 or x in opts.random_artist_categories])
|
||||
artist = random.choice([x for x in shared.artist_db.artists if x.category in allowed_cats])
|
||||
|
||||
return prompt + ", " + artist.name if prompt != '' else artist.name
|
||||
|
||||
|
||||
def visit(x, func, path=""):
|
||||
if hasattr(x, 'children'):
|
||||
for c in x.children:
|
||||
visit(c, func, path)
|
||||
elif x.label is not None:
|
||||
func(path + "/" + str(x.label), x)
|
||||
|
||||
|
||||
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)
|
||||
|
||||
with gr.Row().style(equal_height=False):
|
||||
with gr.Column(variant='panel'):
|
||||
steps = gr.Slider(minimum=1, maximum=150, step=1, label="Sampling Steps", value=20)
|
||||
sampler_index = gr.Radio(label='Sampling method', elem_id="txt2img_sampling", choices=[x.name for x in samplers], value=samplers[0].name, type="index")
|
||||
|
||||
with gr.Row():
|
||||
restore_faces = gr.Checkbox(label='Restore faces', value=False, visible=len(shared.face_restorers) > 1)
|
||||
tiling = gr.Checkbox(label='Tiling', value=False)
|
||||
|
||||
with gr.Row():
|
||||
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)
|
||||
|
||||
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)
|
||||
|
||||
seed = gr.Number(label='Seed', value=-1)
|
||||
|
||||
with gr.Group():
|
||||
custom_inputs = modules.scripts.scripts_txt2img.setup_ui(is_img2img=False)
|
||||
|
||||
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')
|
||||
|
||||
|
||||
with gr.Group():
|
||||
with gr.Row():
|
||||
save = gr.Button('Save')
|
||||
send_to_img2img = gr.Button('Send to img2img')
|
||||
send_to_inpaint = gr.Button('Send to inpaint')
|
||||
send_to_extras = gr.Button('Send to extras')
|
||||
interrupt = gr.Button('Interrupt')
|
||||
|
||||
progressbar = gr.HTML(elem_id="progressbar")
|
||||
|
||||
with gr.Group():
|
||||
html_info = gr.HTML()
|
||||
generation_info = gr.Textbox(visible=False)
|
||||
|
||||
|
||||
txt2img_args = dict(
|
||||
fn=txt2img,
|
||||
_js="submit",
|
||||
inputs=[
|
||||
prompt,
|
||||
negative_prompt,
|
||||
steps,
|
||||
sampler_index,
|
||||
restore_faces,
|
||||
tiling,
|
||||
batch_count,
|
||||
batch_size,
|
||||
cfg_scale,
|
||||
seed,
|
||||
height,
|
||||
width,
|
||||
] + custom_inputs,
|
||||
outputs=[
|
||||
txt2img_gallery,
|
||||
generation_info,
|
||||
html_info
|
||||
]
|
||||
)
|
||||
|
||||
prompt.submit(**txt2img_args)
|
||||
submit.click(**txt2img_args)
|
||||
|
||||
check_progress.click(
|
||||
fn=check_progress_call,
|
||||
show_progress=False,
|
||||
inputs=[],
|
||||
outputs=[progressbar, txt2img_preview, txt2img_preview],
|
||||
)
|
||||
|
||||
|
||||
interrupt.click(
|
||||
fn=lambda: shared.state.interrupt(),
|
||||
inputs=[],
|
||||
outputs=[],
|
||||
)
|
||||
|
||||
save.click(
|
||||
fn=wrap_gradio_call(save_files),
|
||||
inputs=[
|
||||
generation_info,
|
||||
txt2img_gallery,
|
||||
],
|
||||
outputs=[
|
||||
html_info,
|
||||
html_info,
|
||||
html_info,
|
||||
]
|
||||
)
|
||||
|
||||
roll.click(
|
||||
fn=roll_artist,
|
||||
inputs=[
|
||||
prompt,
|
||||
],
|
||||
outputs=[
|
||||
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)
|
||||
|
||||
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)
|
||||
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")
|
||||
|
||||
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")
|
||||
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)
|
||||
|
||||
with gr.Row():
|
||||
inpaint_full_res = gr.Checkbox(label='Inpaint at full resolution', value=False, visible=False)
|
||||
inpainting_mask_invert = gr.Radio(label='Masking mode', choices=['Inpaint masked', 'Inpaint not masked'], value='Inpaint masked', type="index", visible=False)
|
||||
|
||||
with gr.Row():
|
||||
restore_faces = gr.Checkbox(label='Restore faces', value=False, visible=len(shared.face_restorers) > 1)
|
||||
tiling = gr.Checkbox(label='Tiling', value=False)
|
||||
sd_upscale_overlap = gr.Slider(minimum=0, maximum=256, step=16, label='Tile overlap', value=64, visible=False)
|
||||
|
||||
with gr.Row():
|
||||
sd_upscale_upscaler_name = gr.Radio(label='Upscaler', choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name, type="index", visible=False)
|
||||
|
||||
with gr.Row():
|
||||
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)
|
||||
|
||||
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)
|
||||
|
||||
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)
|
||||
|
||||
seed = gr.Number(label='Seed', value=-1)
|
||||
|
||||
with gr.Group():
|
||||
custom_inputs = modules.scripts.scripts_img2img.setup_ui(is_img2img=True)
|
||||
|
||||
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')
|
||||
|
||||
with gr.Group():
|
||||
with gr.Row():
|
||||
save = gr.Button('Save')
|
||||
img2img_send_to_img2img = gr.Button('Send to img2img')
|
||||
img2img_send_to_inpaint = gr.Button('Send to inpaint')
|
||||
img2img_send_to_extras = gr.Button('Send to extras')
|
||||
interrupt = gr.Button('Interrupt')
|
||||
|
||||
progressbar = gr.HTML(elem_id="progressbar")
|
||||
|
||||
with gr.Group():
|
||||
html_info = gr.HTML()
|
||||
generation_info = gr.Textbox(visible=False)
|
||||
|
||||
def apply_mode(mode):
|
||||
is_classic = mode == 0
|
||||
is_inpaint = mode == 1
|
||||
is_loopback = mode == 2
|
||||
is_upscale = mode == 3
|
||||
|
||||
return {
|
||||
init_img: gr_show(not is_inpaint),
|
||||
init_img_with_mask: 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),
|
||||
}
|
||||
|
||||
switch_mode.change(
|
||||
apply_mode,
|
||||
inputs=[switch_mode],
|
||||
outputs=[
|
||||
init_img,
|
||||
init_img_with_mask,
|
||||
mask_blur,
|
||||
inpainting_fill,
|
||||
batch_count,
|
||||
batch_size,
|
||||
sd_upscale_upscaler_name,
|
||||
sd_upscale_overlap,
|
||||
inpaint_full_res,
|
||||
inpainting_mask_invert,
|
||||
]
|
||||
)
|
||||
|
||||
img2img_args = dict(
|
||||
fn=img2img,
|
||||
_js="submit",
|
||||
inputs=[
|
||||
prompt,
|
||||
init_img,
|
||||
init_img_with_mask,
|
||||
steps,
|
||||
sampler_index,
|
||||
mask_blur,
|
||||
inpainting_fill,
|
||||
restore_faces,
|
||||
tiling,
|
||||
switch_mode,
|
||||
batch_count,
|
||||
batch_size,
|
||||
cfg_scale,
|
||||
denoising_strength,
|
||||
seed,
|
||||
height,
|
||||
width,
|
||||
resize_mode,
|
||||
sd_upscale_upscaler_name,
|
||||
sd_upscale_overlap,
|
||||
inpaint_full_res,
|
||||
inpainting_mask_invert,
|
||||
] + custom_inputs,
|
||||
outputs=[
|
||||
img2img_gallery,
|
||||
generation_info,
|
||||
html_info
|
||||
]
|
||||
)
|
||||
|
||||
prompt.submit(**img2img_args)
|
||||
submit.click(**img2img_args)
|
||||
|
||||
check_progress.click(
|
||||
fn=check_progress_call,
|
||||
show_progress=False,
|
||||
inputs=[],
|
||||
outputs=[progressbar, img2img_preview, img2img_preview],
|
||||
)
|
||||
|
||||
interrupt.click(
|
||||
fn=lambda: shared.state.interrupt(),
|
||||
inputs=[],
|
||||
outputs=[],
|
||||
)
|
||||
|
||||
save.click(
|
||||
fn=wrap_gradio_call(save_files),
|
||||
inputs=[
|
||||
generation_info,
|
||||
img2img_gallery,
|
||||
],
|
||||
outputs=[
|
||||
html_info,
|
||||
html_info,
|
||||
html_info,
|
||||
]
|
||||
)
|
||||
|
||||
send_to_img2img.click(
|
||||
fn=lambda x: image_from_url_text(x),
|
||||
_js="extract_image_from_gallery",
|
||||
inputs=[txt2img_gallery],
|
||||
outputs=[init_img],
|
||||
)
|
||||
|
||||
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],
|
||||
)
|
||||
|
||||
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],
|
||||
)
|
||||
|
||||
with gr.Blocks(analytics_enabled=False) as extras_interface:
|
||||
with gr.Row().style(equal_height=False):
|
||||
with gr.Column(variant='panel'):
|
||||
with gr.Group():
|
||||
image = gr.Image(label="Source", source="upload", interactive=True, type="pil")
|
||||
|
||||
upscaling_resize = gr.Slider(minimum=1.0, maximum=4.0, step=0.05, label="Resize", value=2)
|
||||
|
||||
with gr.Group():
|
||||
extras_upscaler_1 = gr.Radio(label='Upscaler 1', choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name, type="index")
|
||||
|
||||
with gr.Group():
|
||||
extras_upscaler_2 = gr.Radio(label='Upscaler 2', choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name, type="index")
|
||||
extras_upscaler_2_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Upscaler 2 visibility", value=1)
|
||||
|
||||
with gr.Group():
|
||||
gfpgan_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="GFPGAN visibility", value=0, interactive=modules.gfpgan_model.have_gfpgan)
|
||||
|
||||
with gr.Group():
|
||||
codeformer_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="CodeFormer visibility", value=0, interactive=modules.codeformer_model.have_codeformer)
|
||||
codeformer_weight = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="CodeFormer weight (0 = maximum effect, 1 = minimum effect)", value=0, interactive=modules.codeformer_model.have_codeformer)
|
||||
|
||||
submit = gr.Button('Generate', elem_id="extras_generate", variant='primary')
|
||||
|
||||
with gr.Column(variant='panel'):
|
||||
result_image = gr.Image(label="Result")
|
||||
html_info_x = gr.HTML()
|
||||
html_info = gr.HTML()
|
||||
|
||||
extras_args = dict(
|
||||
fn=run_extras,
|
||||
inputs=[
|
||||
image,
|
||||
gfpgan_visibility,
|
||||
codeformer_visibility,
|
||||
codeformer_weight,
|
||||
upscaling_resize,
|
||||
extras_upscaler_1,
|
||||
extras_upscaler_2,
|
||||
extras_upscaler_2_visibility,
|
||||
],
|
||||
outputs=[
|
||||
result_image,
|
||||
html_info_x,
|
||||
html_info,
|
||||
]
|
||||
)
|
||||
|
||||
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=[
|
||||
gr.Image(label="Source", source="upload", interactive=True, type="pil"),
|
||||
],
|
||||
outputs=[
|
||||
gr.HTML(),
|
||||
gr.HTML(),
|
||||
gr.HTML(),
|
||||
],
|
||||
allow_flagging="never",
|
||||
analytics_enabled=False,
|
||||
)
|
||||
|
||||
def create_setting_component(key):
|
||||
def fun():
|
||||
return opts.data[key] if key in opts.data else opts.data_labels[key].default
|
||||
|
||||
info = opts.data_labels[key]
|
||||
t = type(info.default)
|
||||
|
||||
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 {}))
|
||||
elif t == str:
|
||||
item = gr.Textbox(label=info.label, value=fun, lines=1)
|
||||
elif t == int:
|
||||
item = gr.Number(label=info.label, value=fun)
|
||||
elif t == bool:
|
||||
item = gr.Checkbox(label=info.label, value=fun)
|
||||
else:
|
||||
raise Exception(f'bad options item type: {str(t)} for key {key}')
|
||||
|
||||
return item
|
||||
|
||||
def run_settings(*args):
|
||||
up = []
|
||||
|
||||
for key, value, comp in zip(opts.data_labels.keys(), args, settings_interface.input_components):
|
||||
opts.data[key] = value
|
||||
up.append(comp.update(value=value))
|
||||
|
||||
opts.save(shared.config_filename)
|
||||
|
||||
return 'Settings saved.', '', ''
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
interfaces = [
|
||||
(txt2img_interface, "txt2img"),
|
||||
(img2img_interface, "img2img"),
|
||||
(extras_interface, "Extras"),
|
||||
(pnginfo_interface, "PNG Info"),
|
||||
(settings_interface, "Settings"),
|
||||
]
|
||||
|
||||
with open(os.path.join(script_path, "style.css"), "r", encoding="utf8") as file:
|
||||
css = file.read()
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
ui_config_file = os.path.join(modules.paths.script_path, 'ui-config.json')
|
||||
ui_settings = {}
|
||||
settings_count = len(ui_settings)
|
||||
error_loading = False
|
||||
|
||||
try:
|
||||
if os.path.exists(ui_config_file):
|
||||
with open(ui_config_file, "r", encoding="utf8") as file:
|
||||
ui_settings = json.load(file)
|
||||
except Exception:
|
||||
error_loading = True
|
||||
print("Error loading settings:", file=sys.stderr)
|
||||
print(traceback.format_exc(), file=sys.stderr)
|
||||
|
||||
def loadsave(path, x):
|
||||
def apply_field(obj, field, condition=None):
|
||||
key = path + "/" + field
|
||||
|
||||
saved_value = ui_settings.get(key, None)
|
||||
if saved_value is None:
|
||||
ui_settings[key] = getattr(obj, field)
|
||||
elif condition is None or condition(saved_value):
|
||||
setattr(obj, field, saved_value)
|
||||
|
||||
if type(x) == gr.Slider:
|
||||
apply_field(x, 'value')
|
||||
apply_field(x, 'minimum')
|
||||
apply_field(x, 'maximum')
|
||||
apply_field(x, 'step')
|
||||
|
||||
if type(x) == gr.Radio:
|
||||
apply_field(x, 'value', lambda val: val in x.choices)
|
||||
|
||||
visit(txt2img_interface, loadsave, "txt2img")
|
||||
visit(img2img_interface, loadsave, "img2img")
|
||||
|
||||
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:
|
||||
json.dump(ui_settings, file, indent=4)
|
||||
|
||||
return demo
|
||||
|
||||
|
||||
with open(os.path.join(script_path, "script.js"), "r", encoding="utf8") as jsfile:
|
||||
javascript = jsfile.read()
|
||||
|
||||
|
||||
def template_response(*args, **kwargs):
|
||||
res = gradio_routes_templates_response(*args, **kwargs)
|
||||
res.body = res.body.replace(b'</head>', f'<script>{javascript}</script></head>'.encode("utf8"))
|
||||
res.init_headers()
|
||||
return res
|
||||
|
||||
|
||||
gradio_routes_templates_response = gradio.routes.templates.TemplateResponse
|
||||
gradio.routes.templates.TemplateResponse = template_response
|
||||
@@ -1,14 +0,0 @@
|
||||
basicsr
|
||||
gfpgan
|
||||
gradio
|
||||
numpy
|
||||
Pillow
|
||||
realesrgan
|
||||
torch
|
||||
transformers
|
||||
omegaconf
|
||||
pytorch_lightning
|
||||
diffusers
|
||||
invisible-watermark
|
||||
git+https://github.com/crowsonkb/k-diffusion.git
|
||||
git+https://github.com/TencentARC/GFPGAN.git
|
||||
@@ -1,10 +0,0 @@
|
||||
basicsr==1.3.5
|
||||
gfpgan
|
||||
gradio==3.2
|
||||
numpy==1.22.0
|
||||
Pillow==9.2.0
|
||||
realesrgan==0.2.5.0
|
||||
torch
|
||||
transformers==4.19.2
|
||||
omegaconf==2.1.1
|
||||
pytorch_lightning==1.7.2
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 513 KiB After Width: | Height: | Size: 865 KiB |
@@ -1,155 +0,0 @@
|
||||
titles = {
|
||||
"Sampling steps": "How many times to imptove the generated image itratively; 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",
|
||||
"DDIM": "Denoising Diffusion Implicit Models - best at inpainting",
|
||||
|
||||
"Batch count": "How many batches of images to create",
|
||||
"Batch size": "How many image to create in a single batch",
|
||||
"CFG Scale": "Classifier Free Guidance Scale - how strongly the image should conform to prompt - lower values produce more creative results",
|
||||
"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.",
|
||||
"Crop and resize": "Resize the image so that entirety of target resolution is filled with the image. Crop parts that stick out.",
|
||||
"Resize and fill": "Resize the image so that entirety of image is inside target resolution. Fill empty space with image's colors.",
|
||||
|
||||
"Mask blur": "How much to blur the mask before processing, in pixels.",
|
||||
"Masked content": "What to put inside the masked area before processing it with Stable Diffusion.",
|
||||
"fill": "fill it with colors of the image",
|
||||
"original": "keep whatever was there originally",
|
||||
"latent noise": "fill it with latent space noise",
|
||||
"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.",
|
||||
|
||||
"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.",
|
||||
|
||||
"X values": "Separate values for X axis using commas.",
|
||||
"Y values": "Separate values for Y axis using commas.",
|
||||
|
||||
"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",
|
||||
|
||||
"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.",
|
||||
|
||||
"Roll": "Add a random artist to the prompt.",
|
||||
}
|
||||
|
||||
function gradioApp(){
|
||||
return document.getElementsByTagName('gradio-app')[0].shadowRoot;
|
||||
}
|
||||
|
||||
global_progressbar = null
|
||||
|
||||
function addTitles(root){
|
||||
root.querySelectorAll('span, button, select').forEach(function(span){
|
||||
tooltip = titles[span.textContent];
|
||||
|
||||
if(!tooltip){
|
||||
tooltip = titles[span.value];
|
||||
}
|
||||
|
||||
if(tooltip){
|
||||
span.title = tooltip;
|
||||
}
|
||||
})
|
||||
|
||||
root.querySelectorAll('select').forEach(function(select){
|
||||
if (select.onchange != null) return;
|
||||
|
||||
select.onchange = function(){
|
||||
select.title = titles[select.value] || "";
|
||||
}
|
||||
})
|
||||
|
||||
progressbar = root.getElementById('progressbar')
|
||||
if(progressbar!= null && progressbar != global_progressbar){
|
||||
global_progressbar = progressbar
|
||||
|
||||
var mutationObserver = new MutationObserver(function(m){
|
||||
txt2img_preview = gradioApp().getElementById('txt2img_preview')
|
||||
txt2img_gallery = gradioApp().getElementById('txt2img_gallery')
|
||||
|
||||
img2img_preview = gradioApp().getElementById('img2img_preview')
|
||||
img2img_gallery = gradioApp().getElementById('img2img_gallery')
|
||||
|
||||
if(txt2img_preview != null && txt2img_gallery != null){
|
||||
txt2img_preview.style.width = txt2img_gallery.clientWidth + "px"
|
||||
txt2img_preview.style.height = txt2img_gallery.clientHeight + "px"
|
||||
}
|
||||
|
||||
if(img2img_preview != null && img2img_gallery != null){
|
||||
img2img_preview.style.width = img2img_gallery.clientWidth + "px"
|
||||
img2img_preview.style.height = img2img_gallery.clientHeight + "px"
|
||||
}
|
||||
|
||||
|
||||
window.setTimeout(requestProgress, 500)
|
||||
});
|
||||
mutationObserver.observe( progressbar, { childList:true, subtree:true })
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
document.addEventListener("DOMContentLoaded", function() {
|
||||
var mutationObserver = new MutationObserver(function(m){
|
||||
addTitles(gradioApp());
|
||||
});
|
||||
mutationObserver.observe( gradioApp(), { childList:true, subtree:true })
|
||||
});
|
||||
|
||||
function selected_gallery_index(){
|
||||
var gr = gradioApp()
|
||||
var buttons = gradioApp().querySelectorAll(".gallery-item")
|
||||
var button = gr.querySelector(".gallery-item.\\!ring-2")
|
||||
|
||||
var result = -1
|
||||
buttons.forEach(function(v, i){ if(v==button) { result = i } })
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
function extract_image_from_gallery(gallery){
|
||||
if(gallery.length == 1){
|
||||
return gallery[0]
|
||||
}
|
||||
|
||||
index = selected_gallery_index()
|
||||
|
||||
if (index < 0 || index >= gallery.length){
|
||||
return []
|
||||
}
|
||||
|
||||
return gallery[index];
|
||||
}
|
||||
|
||||
|
||||
function requestProgress(){
|
||||
btn = gradioApp().getElementById("check_progress");
|
||||
if(btn==null) return;
|
||||
|
||||
btn.click();
|
||||
}
|
||||
|
||||
function submit(){
|
||||
window.setTimeout(requestProgress, 500)
|
||||
|
||||
res = []
|
||||
for(var i=0;i<arguments.length;i++){
|
||||
res.push(arguments[i])
|
||||
}
|
||||
console.log(res)
|
||||
return res
|
||||
}
|
||||
@@ -1,40 +0,0 @@
|
||||
import modules.scripts as scripts
|
||||
import gradio as gr
|
||||
|
||||
from modules.processing import Processed
|
||||
from modules.shared import opts, cmd_opts, state
|
||||
|
||||
|
||||
class Script(scripts.Script):
|
||||
def title(self):
|
||||
return "Custom code"
|
||||
|
||||
|
||||
def show(self, is_img2img):
|
||||
return cmd_opts.allow_code
|
||||
|
||||
def ui(self, is_img2img):
|
||||
code = gr.Textbox(label="Python code", visible=False, lines=1)
|
||||
|
||||
return [code]
|
||||
|
||||
def run(self, p, code):
|
||||
assert cmd_opts.allow_code, '--allow-code option must be enabled'
|
||||
|
||||
display_result_data = [[], -1, ""]
|
||||
|
||||
def display(imgs, s=display_result_data[1], i=display_result_data[2]):
|
||||
display_result_data[0] = imgs
|
||||
display_result_data[1] = s
|
||||
display_result_data[2] = i
|
||||
|
||||
from types import ModuleType
|
||||
compiled = compile(code, '', 'exec')
|
||||
module = ModuleType("testmodule")
|
||||
module.__dict__.update(globals())
|
||||
module.p = p
|
||||
module.display = display
|
||||
exec(compiled, module.__dict__)
|
||||
|
||||
return Processed(p, *display_result_data)
|
||||
|
||||
@@ -1,129 +0,0 @@
|
||||
import math
|
||||
|
||||
import modules.scripts as scripts
|
||||
import gradio as gr
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
from modules import images, processing
|
||||
from modules.processing import Processed, process_images
|
||||
from modules.shared import opts, cmd_opts, state
|
||||
|
||||
|
||||
|
||||
class Script(scripts.Script):
|
||||
def title(self):
|
||||
return "Poor man's outpainting"
|
||||
|
||||
def show(self, is_img2img):
|
||||
return is_img2img
|
||||
|
||||
def ui(self, is_img2img):
|
||||
if not is_img2img:
|
||||
return None
|
||||
|
||||
pixels = gr.Slider(label="Pixels to expand", minimum=8, maximum=128, step=8)
|
||||
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'])
|
||||
|
||||
return [pixels, mask_blur, inpainting_fill, direction]
|
||||
|
||||
def run(self, p, pixels, mask_blur, inpainting_fill, direction):
|
||||
initial_seed = None
|
||||
initial_info = None
|
||||
|
||||
p.mask_blur = mask_blur
|
||||
p.inpainting_fill = inpainting_fill
|
||||
p.inpaint_full_res = False
|
||||
|
||||
left = pixels if "left" in direction else 0
|
||||
right = pixels if "right" in direction else 0
|
||||
up = pixels if "up" in direction else 0
|
||||
down = pixels if "down" in direction else 0
|
||||
|
||||
init_img = p.init_images[0]
|
||||
target_w = math.ceil((init_img.width + left + right) / 64) * 64
|
||||
target_h = math.ceil((init_img.height + up + down) / 64) * 64
|
||||
|
||||
if left > 0:
|
||||
left = left * (target_w - init_img.width) // (left + right)
|
||||
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
|
||||
|
||||
img = Image.new("RGB", (target_w, target_h))
|
||||
img.paste(init_img, (left, up))
|
||||
|
||||
mask = Image.new("L", (img.width, img.height), "white")
|
||||
draw = ImageDraw.Draw(mask)
|
||||
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")
|
||||
|
||||
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")
|
||||
|
||||
processing.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)
|
||||
|
||||
p.n_iter = 1
|
||||
p.batch_size = 1
|
||||
p.do_not_save_grid = True
|
||||
p.do_not_save_samples = True
|
||||
|
||||
work = []
|
||||
work_mask = []
|
||||
work_latent_mask = []
|
||||
work_results = []
|
||||
|
||||
for (_, _, 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):
|
||||
work.append(tiledata[2])
|
||||
work_mask.append(tiledata_mask[2])
|
||||
work_latent_mask.append(tiledata_latent_mask[2])
|
||||
|
||||
batch_count = len(work)
|
||||
print(f"Poor man's outpainting will process a total of {len(work)} images tiled as {len(grid.tiles[0][2])}x{len(grid.tiles)}.")
|
||||
|
||||
state.job_count = batch_count
|
||||
|
||||
for i in range(batch_count):
|
||||
p.init_images = [work[i]]
|
||||
p.image_mask = work_mask[i]
|
||||
p.latent_mask = work_latent_mask[i]
|
||||
|
||||
state.job = f"Batch {i + 1} out of {batch_count}"
|
||||
processed = process_images(p)
|
||||
|
||||
if initial_seed is None:
|
||||
initial_seed = processed.seed
|
||||
initial_info = processed.info
|
||||
|
||||
p.seed = processed.seed + 1
|
||||
work_results += processed.images
|
||||
|
||||
|
||||
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
|
||||
|
||||
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)
|
||||
|
||||
processed = Processed(p, [combined_image], initial_seed, initial_info)
|
||||
|
||||
return processed
|
||||
|
||||
@@ -1,87 +0,0 @@
|
||||
import math
|
||||
from collections import namedtuple
|
||||
from copy import copy
|
||||
import random
|
||||
|
||||
import modules.scripts as scripts
|
||||
import gradio as gr
|
||||
|
||||
from modules import images
|
||||
from modules.processing import process_images, Processed
|
||||
from modules.shared import opts, cmd_opts, state
|
||||
import modules.sd_samplers
|
||||
|
||||
|
||||
def draw_xy_grid(xs, ys, x_label, y_label, cell):
|
||||
res = []
|
||||
|
||||
ver_texts = [[images.GridAnnotation(y_label(y))] for y in ys]
|
||||
hor_texts = [[images.GridAnnotation(x_label(x))] for x in xs]
|
||||
|
||||
first_pocessed = None
|
||||
|
||||
state.job_count = len(xs) * len(ys)
|
||||
|
||||
for iy, y in enumerate(ys):
|
||||
for ix, x in enumerate(xs):
|
||||
state.job = f"{ix + iy * len(xs) + 1} out of {len(xs) * len(ys)}"
|
||||
|
||||
processed = cell(x, y)
|
||||
if first_pocessed is None:
|
||||
first_pocessed = processed
|
||||
|
||||
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)
|
||||
|
||||
first_pocessed.images = [grid]
|
||||
|
||||
return first_pocessed
|
||||
|
||||
|
||||
class Script(scripts.Script):
|
||||
def title(self):
|
||||
return "Prompt matrix"
|
||||
|
||||
def ui(self, is_img2img):
|
||||
put_at_start = gr.Checkbox(label='Put variable parts at start of prompt', value=False)
|
||||
|
||||
return [put_at_start]
|
||||
|
||||
def run(self, p, put_at_start):
|
||||
seed = modules.processing.set_seed(p.seed)
|
||||
|
||||
original_prompt = p.prompt[0] if type(p.prompt) == list else p.prompt
|
||||
|
||||
all_prompts = []
|
||||
prompt_matrix_parts = original_prompt.split("|")
|
||||
combination_count = 2 ** (len(prompt_matrix_parts) - 1)
|
||||
for combination_num in range(combination_count):
|
||||
selected_prompts = [text.strip().strip(',') for n, text in enumerate(prompt_matrix_parts[1:]) if combination_num & (1 << n)]
|
||||
|
||||
if put_at_start:
|
||||
selected_prompts = selected_prompts + [prompt_matrix_parts[0]]
|
||||
else:
|
||||
selected_prompts = [prompt_matrix_parts[0]] + selected_prompts
|
||||
|
||||
all_prompts.append(", ".join(selected_prompts))
|
||||
|
||||
p.n_iter = math.ceil(len(all_prompts) / p.batch_size)
|
||||
p.do_not_save_grid = True
|
||||
|
||||
print(f"Prompt matrix will create {len(all_prompts)} images using a total of {p.n_iter} batches.")
|
||||
|
||||
p.prompt = 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))
|
||||
grid = images.draw_prompt_matrix(grid, p.width, p.height, prompt_matrix_parts)
|
||||
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)
|
||||
|
||||
return processed
|
||||
@@ -1,163 +0,0 @@
|
||||
from collections import namedtuple
|
||||
from copy import copy
|
||||
import random
|
||||
|
||||
import modules.scripts as scripts
|
||||
import gradio as gr
|
||||
|
||||
from modules import images
|
||||
from modules.processing import process_images, Processed
|
||||
from modules.shared import opts, cmd_opts, state
|
||||
import modules.sd_samplers
|
||||
import re
|
||||
|
||||
|
||||
def apply_field(field):
|
||||
def fun(p, x, xs):
|
||||
setattr(p, field, x)
|
||||
|
||||
return fun
|
||||
|
||||
|
||||
def apply_prompt(p, x, xs):
|
||||
p.prompt = p.prompt.replace(xs[0], x)
|
||||
|
||||
|
||||
samplers_dict = {}
|
||||
for i, sampler in enumerate(modules.sd_samplers.samplers):
|
||||
samplers_dict[sampler.name.lower()] = i
|
||||
for alias in sampler.aliases:
|
||||
samplers_dict[alias.lower()] = i
|
||||
|
||||
|
||||
def apply_sampler(p, x, xs):
|
||||
sampler_index = samplers_dict.get(x.lower(), None)
|
||||
if sampler_index is None:
|
||||
raise RuntimeError(f"Unknown sampler: {x}")
|
||||
|
||||
p.sampler_index = sampler_index
|
||||
|
||||
|
||||
def format_value_add_label(p, opt, x):
|
||||
return f"{opt.label}: {x}"
|
||||
|
||||
|
||||
def format_value(p, opt, x):
|
||||
return x
|
||||
|
||||
|
||||
AxisOption = namedtuple("AxisOption", ["label", "type", "apply", "format_value"])
|
||||
AxisOptionImg2Img = namedtuple("AxisOptionImg2Img", ["label", "type", "apply", "format_value"])
|
||||
|
||||
|
||||
axis_options = [
|
||||
AxisOption("Seed", int, apply_field("seed"), 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
|
||||
]
|
||||
|
||||
|
||||
def draw_xy_grid(xs, ys, x_label, y_label, cell):
|
||||
res = []
|
||||
|
||||
ver_texts = [[images.GridAnnotation(y_label(y))] for y in ys]
|
||||
hor_texts = [[images.GridAnnotation(x_label(x))] for x in xs]
|
||||
|
||||
first_pocessed = None
|
||||
|
||||
state.job_count = len(xs) * len(ys)
|
||||
|
||||
for iy, y in enumerate(ys):
|
||||
for ix, x in enumerate(xs):
|
||||
state.job = f"{ix + iy * len(xs) + 1} out of {len(xs) * len(ys)}"
|
||||
|
||||
processed = cell(x, y)
|
||||
if first_pocessed is None:
|
||||
first_pocessed = processed
|
||||
|
||||
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)
|
||||
|
||||
first_pocessed.images = [grid]
|
||||
|
||||
return first_pocessed
|
||||
|
||||
|
||||
re_range = re.compile(r"\s*([+-]?\s*\d+)\s*-\s*([+-]?\s*\d+)(?:\s*\(([+-]\d+)\s*\))?\s*")
|
||||
|
||||
class Script(scripts.Script):
|
||||
def title(self):
|
||||
return "X/Y plot"
|
||||
|
||||
def ui(self, is_img2img):
|
||||
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_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_values = gr.Textbox(label="Y values", visible=False, lines=1)
|
||||
|
||||
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)
|
||||
p.batch_size = 1
|
||||
p.batch_count = 1
|
||||
|
||||
def process_axis(opt, vals):
|
||||
valslist = [x.strip() for x in vals.split(",")]
|
||||
|
||||
if opt.type == int:
|
||||
valslist_ext = []
|
||||
|
||||
for val in valslist:
|
||||
m = re_range.fullmatch(val)
|
||||
if m is not None:
|
||||
|
||||
start = int(m.group(1))
|
||||
end = int(m.group(2))+1
|
||||
step = int(m.group(3)) if m.group(3) is not None else 1
|
||||
|
||||
valslist_ext += list(range(start, end, step))
|
||||
else:
|
||||
valslist_ext.append(val)
|
||||
|
||||
valslist = valslist_ext
|
||||
|
||||
valslist = [opt.type(x) for x in valslist]
|
||||
|
||||
return valslist
|
||||
|
||||
x_opt = axis_options[x_type]
|
||||
xs = process_axis(x_opt, x_values)
|
||||
|
||||
y_opt = axis_options[y_type]
|
||||
ys = process_axis(y_opt, y_values)
|
||||
|
||||
def cell(x, y):
|
||||
pc = copy(p)
|
||||
x_opt.apply(pc, x, xs)
|
||||
y_opt.apply(pc, y, ys)
|
||||
|
||||
return process_images(pc)
|
||||
|
||||
processed = draw_xy_grid(
|
||||
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
|
||||
)
|
||||
|
||||
if opts.grid_save:
|
||||
images.save_image(processed.images[0], p.outpath_grids, "xy_grid", prompt=p.prompt, seed=processed.seed)
|
||||
|
||||
return processed
|
||||
@@ -1,113 +0,0 @@
|
||||
.output-html p {margin: 0 0.5em;}
|
||||
.performance { font-size: 0.85em; color: #444; }
|
||||
|
||||
#txt2img_generate, #img2img_generate{
|
||||
max-width: 13em;
|
||||
}
|
||||
|
||||
#txt2img_roll{
|
||||
min-width: 1em;
|
||||
max-width: 4em;
|
||||
}
|
||||
|
||||
button{
|
||||
align-self: stretch !important;
|
||||
}
|
||||
|
||||
#img2img_prompt, #txt2img_prompt{
|
||||
padding: 0;
|
||||
border: none !important;
|
||||
}
|
||||
|
||||
#img2maskimg .h-60{
|
||||
height: 30rem;
|
||||
}
|
||||
|
||||
.overflow-hidden, .gr-panel{
|
||||
overflow: visible !important;
|
||||
}
|
||||
|
||||
#x_type, #y_type{
|
||||
max-width: 10em;
|
||||
}
|
||||
|
||||
#txt2img_preview, #img2img_preview{
|
||||
position: absolute;
|
||||
width: 320px;
|
||||
left: 0;
|
||||
right: 0;
|
||||
margin-left: auto;
|
||||
margin-right: auto;
|
||||
z-index: 100;
|
||||
}
|
||||
|
||||
#txt2img_preview div.left-0.top-0, #img2img_preview div.left-0.top-0{
|
||||
display: none;
|
||||
}
|
||||
|
||||
fieldset span.text-gray-500, .gr-block.gr-box span.text-gray-500, label.block span{
|
||||
position: absolute;
|
||||
top: -0.6em;
|
||||
line-height: 1.2em;
|
||||
padding: 0 0.5em;
|
||||
margin: 0;
|
||||
|
||||
background-color: white;
|
||||
border-top: 1px solid #eee;
|
||||
border-left: 1px solid #eee;
|
||||
border-right: 1px solid #eee;
|
||||
}
|
||||
|
||||
.dark fieldset span.text-gray-500, .dark .gr-block.gr-box span.text-gray-500, .dark label.block span{
|
||||
background-color: rgb(31, 41, 55);
|
||||
border-top: 1px solid rgb(55 65 81);
|
||||
border-left: 1px solid rgb(55 65 81);
|
||||
border-right: 1px solid rgb(55 65 81);
|
||||
}
|
||||
|
||||
|
||||
.gr-panel div.flex-col div.justify-between label span{
|
||||
margin: 0;
|
||||
}
|
||||
|
||||
.gr-panel div.flex-col div.justify-between div{
|
||||
position: absolute;
|
||||
top: -0.1em;
|
||||
right: 1em;
|
||||
padding: 0 0.5em;
|
||||
}
|
||||
|
||||
input[type="range"]{
|
||||
margin: 0.5em 0 -0.3em 0;
|
||||
}
|
||||
|
||||
#txt2img_sampling label{
|
||||
padding-left: 0.6em;
|
||||
padding-right: 0.6em;
|
||||
}
|
||||
|
||||
|
||||
|
||||
.progressDiv{
|
||||
width: 100%;
|
||||
height: 30px;
|
||||
background: #b4c0cc;
|
||||
border-radius: 8px;
|
||||
}
|
||||
|
||||
.dark .progressDiv{
|
||||
background: #424c5b;
|
||||
}
|
||||
|
||||
.progressDiv .progress{
|
||||
width: 0%;
|
||||
height: 30px;
|
||||
background: #0060df;
|
||||
color: white;
|
||||
font-weight: bold;
|
||||
line-height: 30px;
|
||||
padding: 0 8px 0 0;
|
||||
text-align: right;
|
||||
border-radius: 8px;
|
||||
}
|
||||
|
||||
@@ -1,171 +0,0 @@
|
||||
@echo off
|
||||
|
||||
set PYTHON=python
|
||||
set GIT=git
|
||||
set COMMANDLINE_ARGS=
|
||||
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
|
||||
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
|
||||
|
||||
dir %VENV_DIR%\Scripts\Python.exe >tmp/stdout.txt 2>tmp/stderr.txt
|
||||
if %ERRORLEVEL% == 0 goto :activate_venv
|
||||
|
||||
for /f "delims=" %%i in ('CALL %PYTHON% -c "import sys; print(sys.executable)"') do set PYTHON_FULLNAME="%%i"
|
||||
echo Creating venv in directory %VENV_DIR% using python %PYTHON_FULLNAME%
|
||||
%PYTHON_FULLNAME% -m venv %VENV_DIR% >tmp/stdout.txt 2>tmp/stderr.txt
|
||||
if %ERRORLEVEL% == 0 goto :activate_venv
|
||||
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
|
||||
echo venv %PYTHON%
|
||||
goto :install_torch
|
||||
|
||||
: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%
|
||||
pause
|
||||
exit /b
|
||||
|
||||
:show_stdout_stderr
|
||||
|
||||
echo.
|
||||
echo exit code: %errorlevel%
|
||||
|
||||
for /f %%i in ("tmp\stdout.txt") do set size=%%~zi
|
||||
if %size% equ 0 goto :show_stderr
|
||||
echo.
|
||||
echo stdout:
|
||||
type tmp\stdout.txt
|
||||
|
||||
:show_stderr
|
||||
for /f %%i in ("tmp\stderr.txt") do set size=%%~zi
|
||||
if %size% equ 0 goto :show_stderr
|
||||
echo.
|
||||
echo stderr:
|
||||
type tmp\stderr.txt
|
||||
|
||||
:endofscript
|
||||
|
||||
echo.
|
||||
echo Launch unsuccessful. Exiting.
|
||||
pause
|
||||
Reference in New Issue
Block a user