ldm-finetune
CompVis latent-diffusion finetuned on art (ongo), logo (erlich) and pixel-art (puck) generation.
This repo is modified from glid-3-xl. Aesthetic CLIP embeds are provided by aesthetic-predictor
Quick start (docker required)
The following command will download all weights and run a prediction with your inputs inside a proper docker container.
cog predict r8.im/laion-ai/erlich \
-i prompt="an armchair in the form of an avocado" \
-i negative="" \
-i init_image=@path/to/image \
-i mask=@path/to/mask \
-i guidance_scale=5.0 \
-i steps=100 \
-i batch_size=4 \
-i width=256 \
-i height=256 \
-i init_skip_fraction=0.0 \
-i aesthetic_rating=9 \
-i aesthetic_weight=0.5 \
-i seed=-1 \
-i intermediate_outputs=FalseValid remote image URL's are:
r8.im/laion-ai/erlichr8.im/laion-ai/ongor8.im/laion-ai/puck
Setup
Prerequisites
Please ensure the following dependencies are installed prior to building this repo:
- build-essential
- libopenmpi-dev
- liblzma-dev
- zlib1g-dev
Pytorch
It's a good idea to use a virtual environment or a conda environment.
python3 -m venv .venv
source venv/bin/activate
(venv) $Before installing, you should install pytorch manually by following the instructions at pytorch.org
(venv) $ pip install torch==1.10.1+cu111 torchvision==0.11.2+cu111 torchaudio==0.10.1 -f https://download.pytorch.org/whl/torch_stable.htmlTo check your cuda version, run nvidia-smi.
Install ldm-finetune
You can now install this repo by running pip install -e . in the project directory.
(venv) $ git clone https://github.com/laion-ai/ldm-finetune.git
(venv) $ cd ldm-finetune
(venv) $ pip install -e .
(venv) $ pip install -r requirements.txtCheckpoints
Foundation/Backbone models:
# OpenAI CLIP ViT-L/14
wget -P /root/.cache/clip "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt
### BERT Text Encoder
wget --continue https://dall-3.com/models/glid-3-xl/bert.pt
### kl-f8 VAE backbone
wget --continue https://dall-3.com/models/glid-3-xl/kl-f8.ptLatent Diffusion Stage 2 (diffusion)
There are several stage 2 checkpoints to choose from:
(recommended) jack000 - inpaint.pt
The second finetune from jack000's glid-3-xl adds support for inpainting and can be used for unconditional output as well by setting the inpaint image_embed to zeros. Additionally finetuned to use the CLIP text embed via cross-attention (similar to unCLIP).
wget --continue https://dall-3.com/models/glid-3-xl/inpaint.pt
LAION Finetuning Checkpoints
Laion also finetuned inpaint.pt with the aim of improving logo generation and painting generation.
Erlich
erlich is inpaint.pt finetuned on a dataset collected from LAION-5B named Large Logo Dataset. It consists of roughly 100K images of logos with captions generated via BLIP using aggressive re-ranking and filtering.
wget --continue -O erlich.pt https://huggingface.co/laion/erlich/resolve/main/model/ema_0.9999_120000.ptOngo
Ongo is inpaint.pt finetuned on the Wikiart dataset consisting of about 100K paintings with captions generated via BLIP using aggressive re-ranking and filtering. We also make use of the original captions which contain the author name and the painting title.
wget https://huggingface.co/laion/ongo/resolve/main/ongo.ptLAION - puck.pt
puck has been trained on pixel art. While the underlying kl-f8 encoder seems to struggle somewhat with pixel art, results are still interesting.
wget https://huggingface.co/laion/puck/resolve/main/puck.ptOther
### CompVis - `diffusion.pt`
# The original checkpoint from CompVis trained on `LAION-400M`. May output watermarks.
wget --continue https://dall-3.com/models/glid-3-xl/diffusion.pt
### jack000 - `finetune.pt`
# The first finetune from jack000's [glid-3-xl](https://github.com/jack000/glid-3-xl). Modified to accept a CLIP text embed and finetuned on curated data to help with watermarks. Doesn't support inpainting.
# wget https://dall-3.com/models/glid-3-xl/finetune.pt
Generating images
You can run prediction via python or docker. Currently the docker method is best supported.
Docker/cog
If you have access to a linux machine (or WSL2.0 on Windows 11) with docker installed, you can very easily run models by installing cog:
sudo curl -o /usr/local/bin/cog -L https://github.com/replicate/cog/releases/latest/download/cog_`uname -s`_`uname -m`
sudo chmod +x /usr/local/bin/cogModify the MODEL_PATH in cog_sample.py:
MODEL_PATH = "erlich.pt" # Can be erlich, ongo, puck, etc.Now you can run predictions via docker container using:
cog predict -i prompt="a logo of a fox made of fire"Output will be returned as a base64 string at the end of generation and is also saved locally at current_{batch_idx}.png
Flask API
If you'd like to stand up your own ldm-finetune Flask API, you can run:
cog build -t my_ldm_image
docker run -d -p 5000:5000 --gpus all my_ldm_imagePredictions can then be accessed via HTTP:
curl http://localhost:5000/predictions -X POST \
-H 'Content-Type: application/json' \
-d '{"input": {"prompt": "a logo of a fox made of fire"}}'The output from the API will be a list of base64 strings representing your generations.
Python
You can also use the standalone python scripts from glid-3-xl.
# fast PLMS sampling
(venv) $ python sample.py --model_path erlich.pt --batch_size 6 --num_batches 6 --text "a cyberpunk girl with a scifi neuralink device on her head"
# sample with an init image
(venv) $ python sample.py --init_image picture.jpg --skip_timesteps 10 --model_path ongo.pt --batch_size 6 --num_batches 6 --text "a cyberpunk girl with a scifi neuralink device on her head"Autoedit
Autoedit uses the inpaint model to give the ldm an image prompting function (that works differently from --init_image) It continuously edits random parts of the image to maximize clip score for the text prompt
$ (venv) python autoedit.py \
--model_path inpaint.pt --kl_path kl-f8.pt --bert_path bert.pt \
--text "high quality professional pixel art" --negative "" --prefix autoedit_generations \
--batch_size 16 --width 256 --height 256 --iterations 25 \
--starting_threshold 0.6 --ending_threshold 0.5 \
--starting_radius 5 --ending_radius 0.1 \
--seed -1 --guidance_scale 5.0 --steps 30 \
--aesthetic_rating 9 --aesthetic_weight 0.5 --wandb_name my_autoedit_wandb_artifactFinetuning
See the script below for an example of finetuning your own model from one of the available chekcpoints.
Finetuning Tips/Tricks
- NVIDIA GPU required. You will need an A100 or better to use a batch size of 64. Using less may present stability issues.
- Monitor the
grad_normin the output log. If it ever goes above 1.0 the checkpoint may be ruined due to exploding gradients.- to fix, try reducing the learning rate, decreasing the batch size.
- Train in 32-bit
- Resume with saved optimizer state when possible.
- to fix, try reducing the learning rate, decreasing the batch size.
#!/bin/bash
# Finetune glid-3-xl inpaint.pt on your own webdataset.
# Note: like all one-off scripts, this is likely to become out of date at some point.
# running python scripts/image_train_inpaint.py --help will give you more info.
# model flags
use_fp16=False # TODO can cause more trouble than it's worth.
MODEL_FLAGS="--dropout 0.1 --attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --image_size 32 --learn_sigma False --noise_schedule linear --num_channels 320 --num_heads 8 --num_res_blocks 2 --resblock_updown False --use_fp16 $use_fp16 --use_scale_shift_norm False"
# checkpoint flags
resume_checkpoint="inpaint.pt"
kl_model="kl-f8.pt"
bert_model="bert.pt"
# training flags
epochs=80
shard_size=512
batch_size=32
microbatch=-1
lr=1e-6 # lr=1e-5 seems to be stable. going above 3e-5 is not stable.
ema_rate=0.9999 # TODO you may want to lower this to 0.999, 0.99, 0.95, etc.
random_crop=False
random_flip=False
cache_dir="cache"
image_key="jpg"
caption_key="txt"
data_dir=/my/custom/webdataset/ # TODO set this to a real path
# interval flags
sample_interval=100
log_interval=1
save_interval=2000
CKPT_FLAGS="--kl_model $kl_model --bert_model $bert_model --resume_checkpoint $resume_checkpoint"
INTERVAL_FLAGS="--sample_interval $sample_interval --log_interval $log_interval --save_interval $save_interval"
TRAIN_FLAGS="--epochs $epochs --shard_size $shard_size --batch_size $batch_size --microbatch $microbatch --lr $lr --random_crop $random_crop --random_flip $random_flip --cache_dir $cache_dir --image_key $image_key --caption_key $caption_key --data_dir $data_dir"
COMBINED_FLAGS="$MODEL_FLAGS $CKPT_FLAGS $TRAIN_FLAGS $INTERVAL_FLAGS"
export OPENAI_LOGDIR=./erlich_on_pixel_logs_run6_part2/
export TOKENIZERS_PARALLELISM=false
# TODO comment out a line below to train either on a single GPU or multi-GPU
# single GPU
# python scripts/image_train_inpaint.py $COMBINED_FLAGS
# or multi-GPU
# mpirun -n 8 python scripts/image_train_inpaint.py $COMBINED_FLAGS