This repository contains an additional script, generate.py
, for interacting with both text2img and img2img generation. The script is derived from the optimizedSD scripts, and will feature the same tradeoff of reduced memory consumption with reduced generation speed.
generate.py
contains some additional features, including:
- storing of the prompt and relevant args as PNG metadata
- creating a text file of the prompt and relevant args for each generated image set (same name as the image file, can be used for detecting newly generated images and reading their parameters when automatically maintaining an external image gallery. Especially as this repository does not feature the NSFW filter, images should be manually reviewed before being automatically integrated into a publicly viewable gallery, however.)
- unique, timestamped filenames
- fixing of incorrect input dimensions to multiples of 64
Note that the seed stored in the metadata is a global seed for all images when generating more than one image. To retrieve images of the same seed, both the original seed and the original batch size need to be set.
Prompts are taken as a positional argument, making the simplest way of generating an image using text2img:
python generate.py "painting of a painter painting a painting"
Multiple prompts can also be specified this way, separated by '||'. The batch items (when generating multiple images) will cycle through the prompts:
python generate.py "studio photograph of an apple || masterful psychedelic painting of an apple" -n 4
Generating outputs using img2img is accessed the same way as generating outputs using text2img. When specifying a path to the initial image using either -ii
or --init_img
, generate.py
switches to img2img mode:
python generate.py "a prompt" -ii "./input_image.png"
For an example, see section Image Modification with Stable Diffusion below.
-
This repo is a modified version of the Stable Diffusion repo, optimized to use lesser VRAM than the original by sacrificing on inference speed.
-
It has two python files. 1)
optimized_img2img.py
to generate new image based on a given image and prompt. 2)optimized_txt2img.py
to generate an image based only on a prompt.
This code can now generate a 512x512 image from prior images using less than 4GB of RAM (in less than 20 seconds on RTX 2060)
-
It can generate 512x512 images from a prior image and prompt on a 4GB VRAM GPU in under 20 seconds per image (RTX 2060 in my case).
-
You can use the
--H
&--W
arguments to set the size of the generated image.The maximum size that can fit on 6GB GPU (RTX 2060) is around 576x768. -
For example, the following command will generate 20 512x512 images:
python optimizedSD/optimized_img2img.py --prompt "Austrian alps" --init-img ~/sketch-mountains-input.jpg --strength 0.8 --n_iter 2 --n_samples 5 --H 576 --W 768
-
It can generate 512x512 images from a prompt on a 4GB VRAM GPU in under 25 seconds per image (RTX 2060 in my case).
-
You can use the
--H
&--W
arguments to set the size of the generated image. -
For example, the following command will generate 20 512x512 images:
python optimizedSD/optimized_txt2img.py --prompt "Cyberpunk style image of a Telsa car reflection in rain" --H 512 --W 512 --seed 27 --n_iter 2 --n_samples 10 --ddim_steps 50
-
To get the lowest inference time per image, use the maximum batch size
--n_samples
that can fit on the GPU. Inference time per image will reduce on increasing the batch size, but the required VRAM will also increase. -
If you get a CUDA out of memory error, try reducing the batch size
--n_samples
. If it doesn't work, the other option is to reduce the image width--W
or height--H
or both. -
Mixed Precision is enabled by default. If you don't have a GPU with tensor cores, you can still use mixed precision to run the code using lesser VRAM but the inference time may be larger. And if you feel that the inference is slower, try using the
--precision full
argument to disable it. -
All the modified files are in the optimizedSD folder, so if you have already cloned the original repo, you can just download and copy this folder into the orignal repo instead of cloning the entire repo.
-
You can also clone this repo and follow the same installation steps as the original written below (mainly creating the conda env and placing the weights at the specified location).
-
To achieve this, the stable diffusion model is fragmented into four parts which are sent to the GPU only when needed. After the calculation is done, they are moved back to the CPU. This allows us to run a bigger model on a lower VRAM.
-
The only drawback is higher inference time which is still an order of magnitude faster than inference on CPU.
Stable Diffusion was made possible thanks to a collaboration with Stability AI and Runway and builds upon our previous work:
High-Resolution Image Synthesis with Latent Diffusion Models
Robin Rombach*,
Andreas Blattmann*,
Dominik Lorenz,
Patrick Esser,
Björn Ommer
which is available on GitHub.
Stable Diffusion is a latent text-to-image diffusion
model.
Thanks to a generous compute donation from Stability AI and support from LAION, we were able to train a Latent Diffusion Model on 512x512 images from a subset of the LAION-5B database.
Similar to Google's Imagen,
this model uses a frozen CLIP ViT-L/14 text encoder to condition the model on text prompts.
With its 860M UNet and 123M text encoder, the model is relatively lightweight and runs on a GPU with at least 10GB VRAM.
See this section below and the model card.
A suitable conda environment named ldm
can be created
and activated with:
conda env create -f environment.yaml
conda activate ldm
You can also update an existing latent diffusion environment by running
conda install pytorch torchvision -c pytorch
pip install transformers==4.19.2
pip install -e .
Stable Diffusion v1 refers to a specific configuration of the model architecture that uses a downsampling-factor 8 autoencoder with an 860M UNet and CLIP ViT-L/14 text encoder for the diffusion model. The model was pretrained on 256x256 images and then finetuned on 512x512 images.
*Note: Stable Diffusion v1 is a general text-to-image diffusion model and therefore mirrors biases and (mis-)conceptions that are present in its training data. Details on the training procedure and data, as well as the intended use of the model can be found in the corresponding model card. Research into the safe deployment of general text-to-image models is an ongoing effort. To prevent misuse and harm, we currently provide access to the checkpoints only for academic research purposes upon request. This is an experiment in safe and community-driven publication of a capable and general text-to-image model. We are working on a public release with a more permissive license that also incorporates ethical considerations.*
Request access to Stable Diffusion v1 checkpoints for academic research
We currently provide three checkpoints, sd-v1-1.ckpt
, sd-v1-2.ckpt
and sd-v1-3.ckpt
,
which were trained as follows,
sd-v1-1.ckpt
: 237k steps at resolution256x256
on laion2B-en. 194k steps at resolution512x512
on laion-high-resolution (170M examples from LAION-5B with resolution>= 1024x1024
).sd-v1-2.ckpt
: Resumed fromsd-v1-1.ckpt
. 515k steps at resolution512x512
on "laion-improved-aesthetics" (a subset of laion2B-en, filtered to images with an original size>= 512x512
, estimated aesthetics score> 5.0
, and an estimated watermark probability< 0.5
. The watermark estimate is from the LAION-5B metadata, the aesthetics score is estimated using an improved aesthetics estimator).sd-v1-3.ckpt
: Resumed fromsd-v1-2.ckpt
. 195k steps at resolution512x512
on "laion-improved-aesthetics" and 10% dropping of the text-conditioning to improve classifier-free guidance sampling.
Evaluations with different classifier-free guidance scales (1.5, 2.0, 3.0, 4.0,
5.0, 6.0, 7.0, 8.0) and 50 PLMS sampling
steps show the relative improvements of the checkpoints:
Stable Diffusion is a latent diffusion model conditioned on the (non-pooled) text embeddings of a CLIP ViT-L/14 text encoder.
After obtaining the weights, link them
mkdir -p models/ldm/stable-diffusion-v1/
ln -s <path/to/model.ckpt> models/ldm/stable-diffusion-v1/model.ckpt
and sample with
python scripts/txt2img.py --prompt "a photograph of an astronaut riding a horse" --plms
By default, this uses a guidance scale of --scale 7.5
, Katherine Crowson's implementation of the PLMS sampler,
and renders images of size 512x512 (which it was trained on) in 50 steps. All supported arguments are listed below (type python scripts/txt2img.py --help
).
usage: txt2img.py [-h] [--prompt [PROMPT]] [--outdir [OUTDIR]] [--skip_grid] [--skip_save] [--ddim_steps DDIM_STEPS] [--plms] [--laion400m] [--fixed_code] [--ddim_eta DDIM_ETA] [--n_iter N_ITER] [--H H] [--W W] [--C C] [--f F] [--n_samples N_SAMPLES] [--n_rows N_ROWS]
[--scale SCALE] [--from-file FROM_FILE] [--config CONFIG] [--ckpt CKPT] [--seed SEED] [--precision {full,autocast}]
optional arguments:
-h, --help show this help message and exit
--prompt [PROMPT] the prompt to render
--outdir [OUTDIR] dir to write results to
--skip_grid do not save a grid, only individual samples. Helpful when evaluating lots of samples
--skip_save do not save individual samples. For speed measurements.
--ddim_steps DDIM_STEPS
number of ddim sampling steps
--plms use plms sampling
--laion400m uses the LAION400M model
--fixed_code if enabled, uses the same starting code across samples
--ddim_eta DDIM_ETA ddim eta (eta=0.0 corresponds to deterministic sampling
--n_iter N_ITER sample this often
--H H image height, in pixel space
--W W image width, in pixel space
--C C latent channels
--f F downsampling factor
--n_samples N_SAMPLES
how many samples to produce for each given prompt. A.k.a. batch size
--n_rows N_ROWS rows in the grid (default: n_samples)
--scale SCALE unconditional guidance scale: eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty))
--from-file FROM_FILE
if specified, load prompts from this file
--config CONFIG path to config which constructs model
--ckpt CKPT path to checkpoint of model
--seed SEED the seed (for reproducible sampling)
--precision {full,autocast}
evaluate at this precision
Note: The inference config for all v1 versions is designed to be used with EMA-only checkpoints.
For this reason use_ema=False
is set in the configuration, otherwise the code will try to switch from
non-EMA to EMA weights. If you want to examine the effect of EMA vs no EMA, we provide "full" checkpoints
which contain both types of weights. For these, use_ema=False
will load and use the non-EMA weights.
By using a diffusion-denoising mechanism as first proposed by SDEdit, the model can be used for different tasks such as text-guided image-to-image translation and upscaling. Similar to the txt2img sampling script, we provide a script to perform image modification with Stable Diffusion.
The following describes an example where a rough sketch made in Pinta is converted into a detailed artwork.
python scripts/img2img.py --prompt "A fantasy landscape, trending on artstation" --init-img <path-to-img.jpg> --strength 0.8
Here, strength is a value between 0.0 and 1.0, that controls the amount of noise that is added to the input image. Values that approach 1.0 allow for lots of variations but will also produce images that are not semantically consistent with the input. See the following example.
Input
Outputs
This procedure can, for example, also be used to upscale samples from the base model.
-
Our codebase for the diffusion models builds heavily on OpenAI's ADM codebase and https://github.com/lucidrains/denoising-diffusion-pytorch. Thanks for open-sourcing!
-
The implementation of the transformer encoder is from x-transformers by lucidrains.
@misc{rombach2021highresolution,
title={High-Resolution Image Synthesis with Latent Diffusion Models},
author={Robin Rombach and Andreas Blattmann and Dominik Lorenz and Patrick Esser and Björn Ommer},
year={2021},
eprint={2112.10752},
archivePrefix={arXiv},
primaryClass={cs.CV}
}