├── .dockerignore ├── .gitignore ├── LICENSE ├── README.md ├── aesthetic_clip_embeds ├── rating0.npy ├── rating1.npy ├── rating2.npy ├── rating3.npy ├── rating4.npy ├── rating5.npy ├── rating6.npy ├── rating7.npy ├── rating8.npy └── rating9.npy ├── assets ├── circle_mask.png ├── colorful-glowing-low-poly-logo-of-a-lion.png ├── ongo-painting-of-a-farm-with-flowers.png ├── puck-super-mario-world.png ├── sample.py └── simulacra.txt ├── autoedit.py ├── cog.yaml ├── cog_autoedit.py ├── cog_sample.py ├── dist └── clip_custom │ ├── __init__.py │ ├── bpe_simple_vocab_16e6.txt.gz │ ├── clip.py │ ├── model.py │ └── simple_tokenizer.py ├── encoders ├── modules.py └── x_transformer.py ├── guided_diffusion ├── __init__.py ├── dist_util.py ├── fp16_util.py ├── gaussian_diffusion.py ├── image_text_datasets.py ├── inpaint_util.py ├── logger.py ├── losses.py ├── nn.py ├── predict_util.py ├── resample.py ├── respace.py ├── script_util.py ├── train_util.py └── unet.py ├── ldm ├── models │ ├── autoencoder.py │ └── diffusion │ │ ├── __init__.py │ │ ├── classifier.py │ │ ├── ddim.py │ │ ├── ddpm.py │ │ └── plms.py ├── modules │ ├── attention.py │ ├── diffusionmodules │ │ ├── __init__.py │ │ ├── model.py │ │ ├── openaimodel.py │ │ └── util.py │ ├── distributions │ │ ├── __init__.py │ │ └── distributions.py │ ├── ema.py │ ├── encoders │ │ ├── __init__.py │ │ └── modules.py │ ├── image_degradation │ │ ├── __init__.py │ │ ├── bsrgan.py │ │ ├── bsrgan_light.py │ │ ├── utils │ │ │ └── test.png │ │ └── utils_image.py │ ├── losses │ │ ├── __init__.py │ │ ├── contperceptual.py │ │ └── vqperceptual.py │ └── x_transformer.py └── util.py ├── requirements.txt ├── sample_inpaint.py ├── scripts ├── image_train_inpaint.py └── image_train_latent.py └── setup.py /.dockerignore: -------------------------------------------------------------------------------- 1 | wandb/ 2 | .vscode/ 3 | .idea/ 4 | *.pyc -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .venv/ 2 | .vscode/ 3 | *.pt 4 | *.pyc 5 | auto_outputs/ 6 | auto_outputs_npy/ 7 | wandb/ 8 | latent-diffusion/ 9 | run_autoedit.sh.cog/ 10 | output 11 | output_npy/ 12 | guided_diffusion.egg-info/ 13 | textual.onnx 14 | visual.onnx 15 | *.sh 16 | .cog/ 17 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 OpenAI 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # `ldm-finetune` 2 | 3 | CompVis `latent-diffusion` finetuned on art (ongo), logo (erlich) and pixel-art (puck) generation. 4 | 5 | This repo is modified from [glid-3-xl](https://github.com/jack000/glid-3-xl). Aesthetic CLIP embeds are provided by [aesthetic-predictor](https://github.com/LAION-AI/aesthetic-predictor) 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | - [`ldm-finetune`](#ldm-finetune) 22 | - [Quick start (docker required)](#quick-start-docker-required) 23 | - [Setup](#setup) 24 | - [Prerequisites](#prerequisites) 25 | - [Pytorch](#pytorch) 26 | - [Install ldm-finetune](#install-ldm-finetune) 27 | - [Checkpoints](#checkpoints) 28 | - [Foundation/Backbone models:](#foundationbackbone-models) 29 | - [Latent Diffusion Stage 2 (diffusion)](#latent-diffusion-stage-2-diffusion) 30 | - [(recommended) jack000 - `inpaint.pt`](#recommended-jack000---inpaintpt) 31 | - [LAION Finetuning Checkpoints](#laion-finetuning-checkpoints) 32 | - [Erlich](#erlich) 33 | - [Ongo](#ongo) 34 | - [LAION - `puck.pt`](#laion---puckpt) 35 | - [Other](#other) 36 | - [Generating images](#generating-images) 37 | - [Docker/cog](#dockercog) 38 | - [Flask API](#flask-api) 39 | - [Python](#python) 40 | - [Autoedit](#autoedit) 41 | - [Finetuning](#finetuning) 42 | 43 | ## Quick start (docker required) 44 | 45 | - Install [docker](https://docs.docker.com/get-docker/) 46 | - Install [cog](https://github.com/replicate/cog/) 47 | 48 | The following command will download all weights and run a prediction with your inputs inside a proper docker container. 49 | 50 | ```sh 51 | cog predict r8.im/laion-ai/erlich \ 52 | -i prompt="an armchair in the form of an avocado" \ 53 | -i negative="" \ 54 | -i init_image=@path/to/image \ 55 | -i mask=@path/to/mask \ 56 | -i guidance_scale=5.0 \ 57 | -i steps=100 \ 58 | -i batch_size=4 \ 59 | -i width=256 \ 60 | -i height=256 \ 61 | -i init_skip_fraction=0.0 \ 62 | -i aesthetic_rating=9 \ 63 | -i aesthetic_weight=0.5 \ 64 | -i seed=-1 \ 65 | -i intermediate_outputs=False 66 | ``` 67 | 68 | Valid remote image URL's are: 69 | 70 | - `r8.im/laion-ai/erlich` 71 | - `r8.im/laion-ai/ongo` 72 | - `r8.im/laion-ai/puck` 73 | 74 | ## Setup 75 | 76 | ### Prerequisites 77 | 78 | Please ensure the following dependencies are installed prior to building this repo: 79 | 80 | - build-essential 81 | - libopenmpi-dev 82 | - liblzma-dev 83 | - zlib1g-dev 84 | 85 | 86 | ### Pytorch 87 | 88 | It's a good idea to use a virtual environment or a conda environment. 89 | 90 | ```bash 91 | python3 -m venv .venv 92 | source venv/bin/activate 93 | (venv) $ 94 | ``` 95 | 96 | Before installing, you should install pytorch manually by following the instructions at [pytorch.org](https://pytorch.org/get-started/locally/) 97 | 98 | ```bash 99 | (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.html 100 | ``` 101 | 102 | To check your cuda version, run `nvidia-smi`. 103 | 104 | ### Install ldm-finetune 105 | 106 | You can now install this repo by running `pip install -e .` in the project directory. 107 | 108 | ```bash 109 | (venv) $ git clone https://github.com/laion-ai/ldm-finetune.git 110 | (venv) $ cd ldm-finetune 111 | (venv) $ pip install -e . 112 | (venv) $ pip install -r requirements.txt 113 | ``` 114 | 115 | ## Checkpoints 116 | 117 | ### Foundation/Backbone models: 118 | ```sh 119 | # OpenAI CLIP ViT-L/14 120 | wget -P /root/.cache/clip "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt 121 | 122 | ### BERT Text Encoder 123 | wget --continue https://dall-3.com/models/glid-3-xl/bert.pt 124 | 125 | ### kl-f8 VAE backbone 126 | wget --continue https://dall-3.com/models/glid-3-xl/kl-f8.pt 127 | ``` 128 | 129 | ### Latent Diffusion Stage 2 (diffusion) 130 | There are several stage 2 checkpoints to choose from: 131 | 132 | ### (recommended) jack000 - `inpaint.pt` 133 | 134 | The second finetune from jack000's [glid-3-xl](https://github.com/jack000/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). 135 | 136 | wget --continue https://dall-3.com/models/glid-3-xl/inpaint.pt 137 | 138 | ### LAION Finetuning Checkpoints 139 | 140 | Laion also finetuned `inpaint.pt` with the aim of improving logo generation and painting generation. 141 | 142 | #### Erlich 143 | `erlich` is [inpaint.pt](https://dall-3.com/models/glid-3-xl/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. 144 | 145 | ```sh 146 | wget --continue -O erlich.pt https://huggingface.co/laion/erlich/resolve/main/model/ema_0.9999_120000.pt 147 | ``` 148 | 149 | > ["You know aviato?"](https://www.youtube.com/watch?v=7Q9nQXdzNd0&t=39s) 150 | 151 | #### Ongo 152 | Ongo is [inpaint.pt](https://dall-3.com/models/glid-3-xl/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. 153 | 154 | ```sh 155 | wget https://huggingface.co/laion/ongo/resolve/main/ongo.pt 156 | ``` 157 | 158 | > ["Ongo Gablogian, the art collector. Charmed, I'm sure."](https://www.youtube.com/watch?v=CuMO5q1Syek) 159 | 160 | #### LAION - `puck.pt` 161 | 162 | `puck` has been trained on pixel art. While the underlying kl-f8 encoder seems to struggle somewhat with pixel art, results are still interesting. 163 | 164 | ```sh 165 | wget https://huggingface.co/laion/puck/resolve/main/puck.pt 166 | ``` 167 | 168 | #### Other 169 | 170 | ``` 171 | ### CompVis - `diffusion.pt` 172 | # The original checkpoint from CompVis trained on `LAION-400M`. May output watermarks. 173 | wget --continue https://dall-3.com/models/glid-3-xl/diffusion.pt 174 | 175 | ### jack000 - `finetune.pt` 176 | # 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. 177 | # wget https://dall-3.com/models/glid-3-xl/finetune.pt 178 | ``` 179 | 180 | ## Generating images 181 | 182 | You can run prediction via python or docker. Currently the docker method is best supported. 183 | 184 | ### Docker/cog 185 | 186 | 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`: 187 | 188 | ```sh 189 | sudo curl -o /usr/local/bin/cog -L https://github.com/replicate/cog/releases/latest/download/cog_`uname -s`_`uname -m` 190 | sudo chmod +x /usr/local/bin/cog 191 | ``` 192 | 193 | Modify the `MODEL_PATH` in `cog_sample.py`: 194 | 195 | ```python 196 | MODEL_PATH = "erlich.pt" # Can be erlich, ongo, puck, etc. 197 | ``` 198 | 199 | Now you can run predictions via docker container using: 200 | 201 | ```sh 202 | cog predict -i prompt="a logo of a fox made of fire" 203 | ``` 204 | 205 | Output will be returned as a base64 string at the end of generation and is also saved locally at `current_{batch_idx}.png` 206 | 207 | 208 | ### Flask API 209 | 210 | If you'd like to stand up your own ldm-finetune Flask API, you can run: 211 | 212 | ```sh 213 | cog build -t my_ldm_image 214 | docker run -d -p 5000:5000 --gpus all my_ldm_image 215 | ``` 216 | 217 | Predictions can then be accessed via HTTP: 218 | 219 | ```sh 220 | curl http://localhost:5000/predictions -X POST \ 221 | -H 'Content-Type: application/json' \ 222 | -d '{"input": {"prompt": "a logo of a fox made of fire"}}' 223 | ``` 224 | 225 | The output from the API will be a list of base64 strings representing your generations. 226 | 227 | ### Python 228 | 229 | You can also use the standalone python scripts from `glid-3-xl`. 230 | 231 | ```bash 232 | # fast PLMS sampling 233 | (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" 234 | 235 | # sample with an init image 236 | (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" 237 | ``` 238 | 239 | ### Autoedit 240 | 241 | > Autoedit uses the inpaint model to give the ldm an image prompting function (that works differently from --init_image) 242 | > It continuously edits random parts of the image to maximize clip score for the text prompt 243 | 244 | ```bash 245 | $ (venv) python autoedit.py \ 246 | --model_path inpaint.pt --kl_path kl-f8.pt --bert_path bert.pt \ 247 | --text "high quality professional pixel art" --negative "" --prefix autoedit_generations \ 248 | --batch_size 16 --width 256 --height 256 --iterations 25 \ 249 | --starting_threshold 0.6 --ending_threshold 0.5 \ 250 | --starting_radius 5 --ending_radius 0.1 \ 251 | --seed -1 --guidance_scale 5.0 --steps 30 \ 252 | --aesthetic_rating 9 --aesthetic_weight 0.5 --wandb_name my_autoedit_wandb_artifact 253 | ``` 254 | 255 | ## Finetuning 256 | 257 | See the script below for an example of finetuning your own model from one of the available chekcpoints. 258 | 259 | Finetuning Tips/Tricks 260 | 261 | - NVIDIA GPU required. You will need an A100 or better to use a batch size of 64. Using less may present stability issues. 262 | - Monitor the `grad_norm` in the output log. If it ever goes above 1.0 the checkpoint may be ruined due to exploding gradients. 263 | - to fix, try reducing the learning rate, decreasing the batch size. 264 | - Train in 32-bit 265 | - Resume with saved optimizer state when possible. 266 | 267 | ```bash 268 | #!/bin/bash 269 | # Finetune glid-3-xl inpaint.pt on your own webdataset. 270 | # Note: like all one-off scripts, this is likely to become out of date at some point. 271 | # running python scripts/image_train_inpaint.py --help will give you more info. 272 | 273 | # model flags 274 | use_fp16=False # TODO can cause more trouble than it's worth. 275 | 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" 276 | 277 | # checkpoint flags 278 | resume_checkpoint="inpaint.pt" 279 | kl_model="kl-f8.pt" 280 | bert_model="bert.pt" 281 | 282 | # training flags 283 | epochs=80 284 | shard_size=512 285 | batch_size=32 286 | microbatch=-1 287 | lr=1e-6 # lr=1e-5 seems to be stable. going above 3e-5 is not stable. 288 | ema_rate=0.9999 # TODO you may want to lower this to 0.999, 0.99, 0.95, etc. 289 | random_crop=False 290 | random_flip=False 291 | cache_dir="cache" 292 | image_key="jpg" 293 | caption_key="txt" 294 | data_dir=/my/custom/webdataset/ # TODO set this to a real path 295 | 296 | # interval flags 297 | sample_interval=100 298 | log_interval=1 299 | save_interval=2000 300 | 301 | CKPT_FLAGS="--kl_model $kl_model --bert_model $bert_model --resume_checkpoint $resume_checkpoint" 302 | INTERVAL_FLAGS="--sample_interval $sample_interval --log_interval $log_interval --save_interval $save_interval" 303 | 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" 304 | COMBINED_FLAGS="$MODEL_FLAGS $CKPT_FLAGS $TRAIN_FLAGS $INTERVAL_FLAGS" 305 | export OPENAI_LOGDIR=./erlich_on_pixel_logs_run6_part2/ 306 | export TOKENIZERS_PARALLELISM=false 307 | 308 | # TODO comment out a line below to train either on a single GPU or multi-GPU 309 | # single GPU 310 | # python scripts/image_train_inpaint.py $COMBINED_FLAGS 311 | 312 | # or multi-GPU 313 | # mpirun -n 8 python scripts/image_train_inpaint.py $COMBINED_FLAGS 314 | ``` 315 | -------------------------------------------------------------------------------- /aesthetic_clip_embeds/rating0.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/aesthetic_clip_embeds/rating0.npy -------------------------------------------------------------------------------- /aesthetic_clip_embeds/rating1.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/aesthetic_clip_embeds/rating1.npy -------------------------------------------------------------------------------- /aesthetic_clip_embeds/rating2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/aesthetic_clip_embeds/rating2.npy -------------------------------------------------------------------------------- /aesthetic_clip_embeds/rating3.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/aesthetic_clip_embeds/rating3.npy -------------------------------------------------------------------------------- /aesthetic_clip_embeds/rating4.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/aesthetic_clip_embeds/rating4.npy -------------------------------------------------------------------------------- /aesthetic_clip_embeds/rating5.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/aesthetic_clip_embeds/rating5.npy -------------------------------------------------------------------------------- /aesthetic_clip_embeds/rating6.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/aesthetic_clip_embeds/rating6.npy -------------------------------------------------------------------------------- /aesthetic_clip_embeds/rating7.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/aesthetic_clip_embeds/rating7.npy -------------------------------------------------------------------------------- /aesthetic_clip_embeds/rating8.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/aesthetic_clip_embeds/rating8.npy -------------------------------------------------------------------------------- /aesthetic_clip_embeds/rating9.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/aesthetic_clip_embeds/rating9.npy -------------------------------------------------------------------------------- /assets/circle_mask.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/assets/circle_mask.png -------------------------------------------------------------------------------- /assets/colorful-glowing-low-poly-logo-of-a-lion.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/assets/colorful-glowing-low-poly-logo-of-a-lion.png -------------------------------------------------------------------------------- /assets/ongo-painting-of-a-farm-with-flowers.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/assets/ongo-painting-of-a-farm-with-flowers.png -------------------------------------------------------------------------------- /assets/puck-super-mario-world.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/assets/puck-super-mario-world.png -------------------------------------------------------------------------------- /cog.yaml: -------------------------------------------------------------------------------- 1 | # image: "r8.im/laion-ai/ldm-autoedit" 2 | image: "r8.im/laion-ai/erlich" 3 | # image: "r8.im/laion-ai/erlich-autoedit" 4 | # image: "r8.im/laion-ai/ongo" 5 | # image: "r8.im/laion-ai/ongo-autoedit" 6 | build: 7 | gpu: true 8 | 9 | system_packages: 10 | - "libopenmpi-dev" 11 | - "lzma-dev" 12 | - "liblzma-dev" 13 | - "software-properties-common" 14 | - "build-essential" 15 | - "libnss3-dev" 16 | - "zlib1g-dev" 17 | - "libgdbm-dev" 18 | - "libncurses5-dev" 19 | - "libssl-dev" 20 | - "libffi-dev" 21 | - "libreadline-dev" 22 | - "libsqlite3-dev" 23 | - "libbz2-dev" 24 | 25 | 26 | python_version: "3.8" 27 | 28 | python_packages: 29 | - "axial-positional-embedding==0.2.1" 30 | - "albumentations==0.4.3" 31 | - "blobfile==1.2.9" 32 | - "braceexpand==0.1.7" 33 | - "cachetools==5.0.0" 34 | - "DALL-E==0.1" 35 | - "dalle-pytorch==1.5.2" 36 | - "einops==0.4.1" 37 | - "gpustat==0.6.0" 38 | - "huggingface-hub==0.5.1" 39 | - "imageio==2.9.0" 40 | - "imageio-ffmpeg==0.4.2" 41 | - "omegaconf==2.1.2" 42 | - "pytorch-lightning==1.6" 43 | - "PyYAML==6.0" 44 | - "regex==2022.4.24" 45 | - "rotary-embedding-torch==0.1.5" 46 | - "tokenizers==0.12.1" 47 | - "torch==1.10.1" 48 | - "torchvision==0.11.2" 49 | - "torchmetrics==0.8.0" 50 | - "tqdm==4.64.0" 51 | - "transformers==4.18.0" 52 | - "torch-fidelity==0.3.0" 53 | - "wandb==0.12.17" 54 | 55 | run: 56 | - pip3 install -e "git+https://github.com/CompVis/latent-diffusion.git#egg=latent_diffusion" 57 | - pip3 install -e "git+https://github.com/CompVis/taming-transformers.git#egg=taming_transformers" 58 | - python3 -c 'from transformers import BertTokenizerFast; t = BertTokenizerFast.from_pretrained("bert-base-uncased");' 59 | - wget -P /root/.cache/clip "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt" 60 | 61 | # predict: "cog_autoedit.py:Predictor" 62 | predict: "cog_sample.py:Predictor" 63 | -------------------------------------------------------------------------------- /cog_autoedit.py: -------------------------------------------------------------------------------- 1 | import os 2 | from random import randint 3 | from typing import Iterator, List 4 | 5 | import cog 6 | import torch 7 | 8 | from autoedit import autoedit 9 | from guided_diffusion.predict_util import ( 10 | average_prompt_embed_with_aesthetic_embed, bert_encode_cfg, 11 | load_aesthetic_vit_l_14_embed, load_bert, load_clip_model_and_transform, 12 | load_diffusion_model, load_vae, pack_model_kwargs, prepare_edit) 13 | 14 | os.environ[ 15 | "TOKENIZERS_PARALLELISM" 16 | ] = "false" # required to avoid errors with transformers lib 17 | 18 | MODEL_PATH = "erlich.pt" 19 | KL_PATH = "kl-f8.pt" 20 | BERT_PATH = "bert.pt" 21 | 22 | 23 | class AutoEditOutput(cog.BaseModel): 24 | image: cog.Path 25 | similarity: float 26 | 27 | 28 | class Predictor(cog.BasePredictor): 29 | @torch.inference_mode() 30 | def setup(self): 31 | self.device = torch.device("cuda") 32 | print(f"Loading model from {MODEL_PATH}") 33 | self.model, self.model_params, self.diffusion = load_diffusion_model( 34 | model_path=MODEL_PATH, 35 | steps="27", 36 | use_fp16=False, 37 | device=self.device, 38 | ) 39 | print(f"Loading vae") 40 | self.ldm = load_vae(kl_path=KL_PATH, device=self.device) 41 | self.ldm = self.ldm 42 | print(f"Loading CLIP") 43 | self.clip_model, self.clip_preprocess = load_clip_model_and_transform(self.device) 44 | print(f"Loading BERT") 45 | self.bert = load_bert(BERT_PATH, self.device) 46 | self.bert = self.bert 47 | 48 | @torch.inference_mode() 49 | def predict( 50 | self, 51 | text: str = cog.Input( 52 | default="", 53 | description="(optional) Text to use for the model's prediction.", 54 | ), 55 | edit: str = cog.Input( 56 | default="", 57 | description="path to the image you want to edit", 58 | ), 59 | negative: str = cog.Input( 60 | default="", 61 | description="(optional) Negate the model's prediction for this text from the model's prediction for the target text.", 62 | ), 63 | aesthetic_rating: int = cog.Input( 64 | description="Number between 0 and 9 representing the aesthetic rating. Will initialize the prompt CLIP embed with the respective aesthetic embed.", 65 | default=9, 66 | ge=0, 67 | le=9, 68 | ), 69 | aesthetic_weight: float = cog.Input( 70 | description="Weight of the aesthetic embedding in the average prompt embedding.", 71 | default=0.5, 72 | ge=0, 73 | le=1, 74 | ), 75 | batch_size: int = cog.Input( 76 | default=1, description="Batch size.", choices=[1, 2, 3, 4, 6, 8] 77 | ), 78 | width: int = cog.Input( 79 | default=256, 80 | description="Target width", 81 | choices=[128, 192, 256, 320, 384], 82 | ), 83 | height: int = cog.Input( 84 | default=256, 85 | description="Target height", 86 | choices=[128, 192, 256, 320, 384], 87 | ), 88 | iterations: int = cog.Input( 89 | default=25, 90 | description="Number of iterations to run the model for.", 91 | ge=25, 92 | ), 93 | starting_radius: float = cog.Input( 94 | default=5.0, 95 | description="size of noise blur at the start of editing (larger = coarser changes)", 96 | ge=0.1, 97 | ), 98 | ending_radius: float = cog.Input( 99 | default=0.1, 100 | description="size of noise blur at the end of editing (smaller = editing fine details)", 101 | ge=0.1, 102 | le=5.0, 103 | ), 104 | starting_threshold: float = cog.Input( 105 | default=0.6, 106 | description="how much of the image to replace at the start of editing (1 = inpaint the entire image)", 107 | ge=0.05, 108 | le=1.0, 109 | ), 110 | ending_threshold: float = cog.Input( 111 | default=0.5, 112 | description="how much of the image to replace at the end of editing", 113 | ge=0.1, 114 | le=1.0, 115 | ), 116 | guidance_scale: float = cog.Input( 117 | default=5.0, 118 | description="Controls how much the image should look like the prompt", 119 | ge=-10.0, 120 | le=100.0, 121 | ), 122 | seed: int = cog.Input( 123 | default=-1, 124 | description="(optional) Seed for the random number generator.", 125 | ge=-1, 126 | ), 127 | ) -> Iterator[List[cog.Path]]: 128 | if seed > 0: 129 | torch.manual_seed(seed) 130 | else: 131 | seed = randint(0, 2**32) 132 | torch.manual_seed(seed) 133 | print(f"Using seed {seed}") 134 | print(f"Running simulation for {text}") 135 | # Create new run and table for each prompt. 136 | prefix = ( 137 | text.replace(" ", "_").replace(",", "_").replace(".", "_").replace("'", "_") 138 | ) 139 | prefix = prefix[:255] 140 | 141 | # Text Setup 142 | print(f"Encoding text embeddings with {text} dimensions") 143 | text_emb, text_blank = bert_encode_cfg( 144 | text, negative, batch_size, self.device, self.bert 145 | ) 146 | text_emb_clip_blank, text_emb_clip, text_emb_norm = clip_encode_cfg( 147 | clip_model=self.clip_model, 148 | text=text, 149 | negative=negative, 150 | batch_size=batch_size, 151 | device=self.device, 152 | ) 153 | print( 154 | f"Using aesthetic embedding {aesthetic_rating} with weight {aesthetic_weight}" 155 | ) 156 | text_emb_clip_aesthetic = load_aesthetic_vit_l_14_embed( 157 | rating=aesthetic_rating 158 | ).to(self.device) 159 | text_emb_clip = average_prompt_embed_with_aesthetic_embed( 160 | text_emb_clip, text_emb_clip_aesthetic, aesthetic_weight 161 | ) 162 | # Image Setup 163 | image_embed = None 164 | if edit: 165 | image_embed = prepare_edit( 166 | self.ldm, edit, batch_size, width, height, self.device 167 | ) 168 | print("Image embedding shape:", image_embed.shape) 169 | elif self.model_params["image_condition"]: 170 | print( 171 | "Using inpaint model but no image is provided. Initializing with zeros." 172 | ) 173 | image_embed = torch.zeros( 174 | batch_size * 2, 4, height // 8, width // 8, device=self.device 175 | ) 176 | 177 | # Prepare inputs 178 | kwargs = pack_model_kwargs( 179 | text_emb=text_emb, 180 | text_blank=text_blank, 181 | text_emb_clip=text_emb_clip, 182 | text_emb_clip_blank=text_emb_clip_blank, 183 | image_embed=image_embed, 184 | model_params=self.model_params, 185 | ) 186 | 187 | for results in autoedit( 188 | model=self.model, 189 | diffusion=self.diffusion, 190 | ldm=self.ldm, 191 | text_emb_norm=text_emb_norm, 192 | clip_model=self.clip_model, 193 | clip_preprocess=self.clip_preprocess, 194 | model_kwargs=kwargs, 195 | batch_size=batch_size, 196 | prefix=prefix, 197 | device=self.device, 198 | guidance_scale=guidance_scale, 199 | width=width, 200 | height=height, 201 | num_mutations=iterations, 202 | starting_radius=starting_radius, 203 | ending_radius=ending_radius, 204 | starting_threshold=starting_threshold, 205 | ending_threshold=ending_threshold, 206 | ): 207 | outputs = [] 208 | for result in results: 209 | decoded_image_path, _, _, similarity = result 210 | # outputs.append(AutoEditOutput(image=cog.Path(str(decoded_image_path)), similarity=similarity)) 211 | outputs.append(cog.Path(str(decoded_image_path))) 212 | yield outputs -------------------------------------------------------------------------------- /cog_sample.py: -------------------------------------------------------------------------------- 1 | import os 2 | import typing 3 | 4 | import cog 5 | import torch 6 | 7 | from guided_diffusion.inpaint_util import (prepare_inpaint_models, 8 | sample_inpaint) 9 | 10 | os.environ[ 11 | "TOKENIZERS_PARALLELISM" 12 | ] = "false" # required to avoid errors with transformers lib 13 | 14 | inpaint_model_path = "inpaint.pt" 15 | 16 | class Predictor(cog.BasePredictor): 17 | @torch.inference_mode() 18 | def setup(self): 19 | self.device = "cuda" if torch.cuda.is_available() else "cpu" 20 | self.use_fp16 = True 21 | self.inpaint_models = prepare_inpaint_models( 22 | inpaint_model_path=inpaint_model_path, 23 | device=self.device, 24 | use_fp16=self.use_fp16, 25 | ) 26 | 27 | @torch.inference_mode() 28 | def predict( 29 | self, 30 | prompt: str = cog.Input(description="Your text prompt.", default=""), 31 | negative: str = cog.Input( 32 | default="", 33 | description="(optional) Negate the model's prediction for this text from the model's prediction for the target text.", 34 | ), 35 | init_image: cog.Path = cog.Input( 36 | default=None, 37 | description="(optional) Initial image to use for the model's prediction. If provided alongside a mask, the image will be inpainted instead.", 38 | ), 39 | mask: cog.Path = cog.Input( 40 | default=None, 41 | description="a mask image for inpainting an init_image. white pixels = keep, black pixels = discard. resized to width = image width/8, height = image height/8", 42 | ), 43 | guidance_scale: float = cog.Input( 44 | default=5.0, 45 | description="Classifier-free guidance scale. Higher values will result in more guidance toward caption, with diminishing returns. Try values between 1.0 and 40.0. In general, going above 5.0 will introduce some artifacting.", 46 | le=100.0, 47 | ge=-20.0, 48 | ), 49 | steps: int = cog.Input( 50 | default=50, 51 | description="Number of diffusion steps to run. Due to PLMS sampling, using more than 100 steps is unnecessary and may simply produce the exact same output.", 52 | le=250, 53 | ge=15, 54 | ), 55 | batch_size: int = cog.Input( 56 | default=4, 57 | description="Batch size. (higher = slower)", 58 | ge=1, 59 | le=16, 60 | ), 61 | width: int = cog.Input( 62 | default=256, 63 | description="Target width", 64 | choices=[128, 192, 256, 320, 384], 65 | ), 66 | height: int = cog.Input( 67 | default=256, 68 | description="Target height", 69 | choices=[128, 192, 256, 320, 384], 70 | ), 71 | init_skip_fraction: float = cog.Input( 72 | default=0.0, 73 | description="Fraction of sampling steps to skip when using an init image. Defaults to 0.0 if init_image is not specified and 0.5 if init_image is specified.", 74 | ge=0.0, 75 | le=1.0, 76 | ), 77 | aesthetic_rating: int = cog.Input( 78 | description="Aesthetic rating (1-9) - embed to use.", default=9 79 | ), 80 | aesthetic_weight: float = cog.Input( 81 | description="Aesthetic weight (0-1). How much to guide towards the aesthetic embed vs the prompt embed.", 82 | default=0.5, 83 | ), 84 | seed: int = cog.Input( 85 | default=-1, 86 | description="Seed for random number generator. If -1, a random seed will be chosen.", 87 | ge=-1, 88 | le=(2**32 - 1), 89 | ), 90 | intermediate_outputs: bool = cog.Input( 91 | default=False, 92 | description="Whether to return intermediate outputs. Enable to visualize the diffusion process and/or debug the model. May slow down inference.", 93 | ), 94 | ) -> typing.Iterator[typing.List[cog.Path]]: 95 | for current_predictions in sample_inpaint( 96 | prompt=prompt, 97 | negative=negative, 98 | init_image=str(init_image) if init_image else None, 99 | mask=str(mask) if mask else None, 100 | steps=steps, 101 | init_skip_fraction=init_skip_fraction, 102 | width=width, 103 | height=height, 104 | batch_size=batch_size, 105 | intermediate_outputs=intermediate_outputs, 106 | guidance_scale=guidance_scale, 107 | aesthetic_rating=aesthetic_rating, 108 | aesthetic_weight=aesthetic_weight, 109 | device=self.device, 110 | use_fp16=self.use_fp16, 111 | seed=seed, 112 | loaded_models=self.inpaint_models, 113 | ): 114 | yield [cog.Path(p) for p in current_predictions] 115 | -------------------------------------------------------------------------------- /dist/clip_custom/__init__.py: -------------------------------------------------------------------------------- 1 | from .clip import * 2 | -------------------------------------------------------------------------------- /dist/clip_custom/bpe_simple_vocab_16e6.txt.gz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/dist/clip_custom/bpe_simple_vocab_16e6.txt.gz -------------------------------------------------------------------------------- /dist/clip_custom/clip.py: -------------------------------------------------------------------------------- 1 | import hashlib 2 | import os 3 | import urllib 4 | import warnings 5 | from typing import Any, Union, List 6 | from pkg_resources import packaging 7 | 8 | import torch 9 | from PIL import Image 10 | from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize 11 | from tqdm import tqdm 12 | 13 | from .model import build_model 14 | from .simple_tokenizer import SimpleTokenizer as _Tokenizer 15 | 16 | try: 17 | from torchvision.transforms import InterpolationMode 18 | BICUBIC = InterpolationMode.BICUBIC 19 | except ImportError: 20 | BICUBIC = Image.BICUBIC 21 | 22 | 23 | if packaging.version.parse(torch.__version__) < packaging.version.parse("1.7.1"): 24 | warnings.warn("PyTorch version 1.7.1 or higher is recommended") 25 | 26 | 27 | __all__ = ["available_models", "load", "tokenize"] 28 | _tokenizer = _Tokenizer() 29 | 30 | _MODELS = { 31 | "RN50": "https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt", 32 | "RN101": "https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt", 33 | "RN50x4": "https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt", 34 | "RN50x16": "https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt", 35 | "RN50x64": "https://openaipublic.azureedge.net/clip/models/be1cfb55d75a9666199fb2206c106743da0f6468c9d327f3e0d0a543a9919d9c/RN50x64.pt", 36 | "ViT-B/32": "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt", 37 | "ViT-B/16": "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt", 38 | "ViT-L/14": "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt", 39 | } 40 | 41 | 42 | def _download(url: str, root: str): 43 | os.makedirs(root, exist_ok=True) 44 | filename = os.path.basename(url) 45 | 46 | expected_sha256 = url.split("/")[-2] 47 | download_target = os.path.join(root, filename) 48 | 49 | if os.path.exists(download_target) and not os.path.isfile(download_target): 50 | raise RuntimeError(f"{download_target} exists and is not a regular file") 51 | 52 | if os.path.isfile(download_target): 53 | if hashlib.sha256(open(download_target, "rb").read()).hexdigest() == expected_sha256: 54 | return download_target 55 | else: 56 | warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file") 57 | 58 | with urllib.request.urlopen(url) as source, open(download_target, "wb") as output: 59 | with tqdm(total=int(source.info().get("Content-Length")), ncols=80, unit='iB', unit_scale=True, unit_divisor=1024) as loop: 60 | while True: 61 | buffer = source.read(8192) 62 | if not buffer: 63 | break 64 | 65 | output.write(buffer) 66 | loop.update(len(buffer)) 67 | 68 | if hashlib.sha256(open(download_target, "rb").read()).hexdigest() != expected_sha256: 69 | raise RuntimeError(f"Model has been downloaded but the SHA256 checksum does not not match") 70 | 71 | return download_target 72 | 73 | 74 | def _convert_image_to_rgb(image): 75 | return image.convert("RGB") 76 | 77 | 78 | def _transform(n_px): 79 | return Compose([ 80 | Resize(n_px, interpolation=BICUBIC), 81 | CenterCrop(n_px), 82 | _convert_image_to_rgb, 83 | ToTensor(), 84 | Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)), 85 | ]) 86 | 87 | 88 | def available_models() -> List[str]: 89 | """Returns the names of available CLIP models""" 90 | return list(_MODELS.keys()) 91 | 92 | 93 | def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", jit: bool = False, download_root: str = None): 94 | """Load a CLIP model 95 | 96 | Parameters 97 | ---------- 98 | name : str 99 | A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict 100 | 101 | device : Union[str, torch.device] 102 | The device to put the loaded model 103 | 104 | jit : bool 105 | Whether to load the optimized JIT model or more hackable non-JIT model (default). 106 | 107 | download_root: str 108 | path to download the model files; by default, it uses "~/.cache/clip" 109 | 110 | Returns 111 | ------- 112 | model : torch.nn.Module 113 | The CLIP model 114 | 115 | preprocess : Callable[[PIL.Image], torch.Tensor] 116 | A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input 117 | """ 118 | if name in _MODELS: 119 | model_path = _download(_MODELS[name], download_root or os.path.expanduser("~/.cache/clip")) 120 | elif os.path.isfile(name): 121 | model_path = name 122 | else: 123 | raise RuntimeError(f"Model {name} not found; available models = {available_models()}") 124 | 125 | try: 126 | # loading JIT archive 127 | model = torch.jit.load(model_path, map_location=device if jit else "cpu").eval() 128 | state_dict = None 129 | except RuntimeError: 130 | # loading saved state dict 131 | if jit: 132 | warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead") 133 | jit = False 134 | state_dict = torch.load(model_path, map_location="cpu") 135 | 136 | if not jit: 137 | model = build_model(state_dict or model.state_dict()).to(device) 138 | if str(device) == "cpu": 139 | model.float() 140 | return model, _transform(model.visual.input_resolution) 141 | 142 | # patch the device names 143 | device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[]) 144 | device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1] 145 | 146 | def patch_device(module): 147 | try: 148 | graphs = [module.graph] if hasattr(module, "graph") else [] 149 | except RuntimeError: 150 | graphs = [] 151 | 152 | if hasattr(module, "forward1"): 153 | graphs.append(module.forward1.graph) 154 | 155 | for graph in graphs: 156 | for node in graph.findAllNodes("prim::Constant"): 157 | if "value" in node.attributeNames() and str(node["value"]).startswith("cuda"): 158 | node.copyAttributes(device_node) 159 | 160 | model.apply(patch_device) 161 | patch_device(model.encode_image) 162 | patch_device(model.encode_text) 163 | 164 | # patch dtype to float32 on CPU 165 | if str(device) == "cpu": 166 | float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[]) 167 | float_input = list(float_holder.graph.findNode("aten::to").inputs())[1] 168 | float_node = float_input.node() 169 | 170 | def patch_float(module): 171 | try: 172 | graphs = [module.graph] if hasattr(module, "graph") else [] 173 | except RuntimeError: 174 | graphs = [] 175 | 176 | if hasattr(module, "forward1"): 177 | graphs.append(module.forward1.graph) 178 | 179 | for graph in graphs: 180 | for node in graph.findAllNodes("aten::to"): 181 | inputs = list(node.inputs()) 182 | for i in [1, 2]: # dtype can be the second or third argument to aten::to() 183 | if inputs[i].node()["value"] == 5: 184 | inputs[i].node().copyAttributes(float_node) 185 | 186 | model.apply(patch_float) 187 | patch_float(model.encode_image) 188 | patch_float(model.encode_text) 189 | 190 | model.float() 191 | 192 | return model, _transform(model.input_resolution.item()) 193 | 194 | 195 | def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> torch.LongTensor: 196 | """ 197 | Returns the tokenized representation of given input string(s) 198 | 199 | Parameters 200 | ---------- 201 | texts : Union[str, List[str]] 202 | An input string or a list of input strings to tokenize 203 | 204 | context_length : int 205 | The context length to use; all CLIP models use 77 as the context length 206 | 207 | truncate: bool 208 | Whether to truncate the text in case its encoding is longer than the context length 209 | 210 | Returns 211 | ------- 212 | A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length] 213 | """ 214 | if isinstance(texts, str): 215 | texts = [texts] 216 | 217 | sot_token = _tokenizer.encoder["<|startoftext|>"] 218 | eot_token = _tokenizer.encoder["<|endoftext|>"] 219 | all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts] 220 | result = torch.zeros(len(all_tokens), context_length, dtype=torch.long) 221 | 222 | for i, tokens in enumerate(all_tokens): 223 | if len(tokens) > context_length: 224 | if truncate: 225 | tokens = tokens[:context_length] 226 | tokens[-1] = eot_token 227 | else: 228 | raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}") 229 | result[i, :len(tokens)] = torch.tensor(tokens) 230 | 231 | return result 232 | -------------------------------------------------------------------------------- /dist/clip_custom/simple_tokenizer.py: -------------------------------------------------------------------------------- 1 | import gzip 2 | import html 3 | import os 4 | from functools import lru_cache 5 | 6 | import ftfy 7 | import regex as re 8 | 9 | 10 | @lru_cache() 11 | def default_bpe(): 12 | return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz") 13 | 14 | 15 | @lru_cache() 16 | def bytes_to_unicode(): 17 | """ 18 | Returns list of utf-8 byte and a corresponding list of unicode strings. 19 | The reversible bpe codes work on unicode strings. 20 | This means you need a large # of unicode characters in your vocab if you want to avoid UNKs. 21 | When you're at something like a 10B token dataset you end up needing around 5K for decent coverage. 22 | This is a signficant percentage of your normal, say, 32K bpe vocab. 23 | To avoid that, we want lookup tables between utf-8 bytes and unicode strings. 24 | And avoids mapping to whitespace/control characters the bpe code barfs on. 25 | """ 26 | bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1)) 27 | cs = bs[:] 28 | n = 0 29 | for b in range(2**8): 30 | if b not in bs: 31 | bs.append(b) 32 | cs.append(2**8+n) 33 | n += 1 34 | cs = [chr(n) for n in cs] 35 | return dict(zip(bs, cs)) 36 | 37 | 38 | def get_pairs(word): 39 | """Return set of symbol pairs in a word. 40 | Word is represented as tuple of symbols (symbols being variable-length strings). 41 | """ 42 | pairs = set() 43 | prev_char = word[0] 44 | for char in word[1:]: 45 | pairs.add((prev_char, char)) 46 | prev_char = char 47 | return pairs 48 | 49 | 50 | def basic_clean(text): 51 | text = ftfy.fix_text(text) 52 | text = html.unescape(html.unescape(text)) 53 | return text.strip() 54 | 55 | 56 | def whitespace_clean(text): 57 | text = re.sub(r'\s+', ' ', text) 58 | text = text.strip() 59 | return text 60 | 61 | 62 | class SimpleTokenizer(object): 63 | def __init__(self, bpe_path: str = default_bpe()): 64 | self.byte_encoder = bytes_to_unicode() 65 | self.byte_decoder = {v: k for k, v in self.byte_encoder.items()} 66 | merges = gzip.open(bpe_path).read().decode("utf-8").split('\n') 67 | merges = merges[1:49152-256-2+1] 68 | merges = [tuple(merge.split()) for merge in merges] 69 | vocab = list(bytes_to_unicode().values()) 70 | vocab = vocab + [v+'' for v in vocab] 71 | for merge in merges: 72 | vocab.append(''.join(merge)) 73 | vocab.extend(['<|startoftext|>', '<|endoftext|>']) 74 | self.encoder = dict(zip(vocab, range(len(vocab)))) 75 | self.decoder = {v: k for k, v in self.encoder.items()} 76 | self.bpe_ranks = dict(zip(merges, range(len(merges)))) 77 | self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'} 78 | self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE) 79 | 80 | def bpe(self, token): 81 | if token in self.cache: 82 | return self.cache[token] 83 | word = tuple(token[:-1]) + ( token[-1] + '',) 84 | pairs = get_pairs(word) 85 | 86 | if not pairs: 87 | return token+'' 88 | 89 | while True: 90 | bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf'))) 91 | if bigram not in self.bpe_ranks: 92 | break 93 | first, second = bigram 94 | new_word = [] 95 | i = 0 96 | while i < len(word): 97 | try: 98 | j = word.index(first, i) 99 | new_word.extend(word[i:j]) 100 | i = j 101 | except: 102 | new_word.extend(word[i:]) 103 | break 104 | 105 | if word[i] == first and i < len(word)-1 and word[i+1] == second: 106 | new_word.append(first+second) 107 | i += 2 108 | else: 109 | new_word.append(word[i]) 110 | i += 1 111 | new_word = tuple(new_word) 112 | word = new_word 113 | if len(word) == 1: 114 | break 115 | else: 116 | pairs = get_pairs(word) 117 | word = ' '.join(word) 118 | self.cache[token] = word 119 | return word 120 | 121 | def encode(self, text): 122 | bpe_tokens = [] 123 | text = whitespace_clean(basic_clean(text)).lower() 124 | for token in re.findall(self.pat, text): 125 | token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8')) 126 | bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' ')) 127 | return bpe_tokens 128 | 129 | def decode(self, tokens): 130 | text = ''.join([self.decoder[token] for token in tokens]) 131 | text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('', ' ') 132 | return text 133 | -------------------------------------------------------------------------------- /encoders/modules.py: -------------------------------------------------------------------------------- 1 | from functools import partial 2 | 3 | import torch 4 | import torch.nn as nn 5 | 6 | from encoders.x_transformer import ( # TODO: can we directly rely on lucidrains code and simply add this as a reuirement? --> test 7 | Encoder, TransformerWrapper) 8 | 9 | 10 | class AbstractEncoder(nn.Module): 11 | def __init__(self): 12 | super().__init__() 13 | 14 | def encode(self, *args, **kwargs): 15 | raise NotImplementedError 16 | 17 | 18 | 19 | class ClassEmbedder(nn.Module): 20 | def __init__(self, embed_dim, n_classes=1000, key='class'): 21 | super().__init__() 22 | self.key = key 23 | self.embedding = nn.Embedding(n_classes, embed_dim) 24 | 25 | def forward(self, batch, key=None): 26 | if key is None: 27 | key = self.key 28 | # this is for use in crossattn 29 | c = batch[key][:, None] 30 | c = self.embedding(c) 31 | return c 32 | 33 | 34 | class TransformerEmbedder(AbstractEncoder): 35 | """Some transformer encoder layers""" 36 | def __init__(self, n_embed, n_layer, vocab_size, max_seq_len=77, device="cuda"): 37 | super().__init__() 38 | self.device = device 39 | self.transformer = TransformerWrapper(num_tokens=vocab_size, max_seq_len=max_seq_len, 40 | attn_layers=Encoder(dim=n_embed, depth=n_layer)) 41 | 42 | def forward(self, tokens): 43 | tokens = tokens.to(self.device) # meh 44 | z = self.transformer(tokens, return_embeddings=True) 45 | return z 46 | 47 | def encode(self, x): 48 | return self(x) 49 | 50 | 51 | class BERTTokenizer(AbstractEncoder): 52 | """ Uses a pretrained BERT tokenizer by huggingface. Vocab size: 30522 (?)""" 53 | def __init__(self, device="cuda", vq_interface=True, max_length=77): 54 | super().__init__() 55 | from transformers import \ 56 | BertTokenizerFast # TODO: add to reuquirements 57 | self.tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased") 58 | self.device = device 59 | self.vq_interface = vq_interface 60 | self.max_length = max_length 61 | 62 | def forward(self, text): 63 | batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True, 64 | return_overflowing_tokens=False, padding="max_length", return_tensors="pt") 65 | tokens = batch_encoding["input_ids"].to(self.device) 66 | return tokens 67 | 68 | @torch.no_grad() 69 | def encode(self, text): 70 | tokens = self(text) 71 | if not self.vq_interface: 72 | return tokens 73 | return None, None, [None, None, tokens] 74 | 75 | def decode(self, text): 76 | return text 77 | 78 | 79 | class BERTEmbedder(AbstractEncoder): 80 | """Uses the BERT tokenizr model and add some transformer encoder layers""" 81 | def __init__(self, n_embed, n_layer, vocab_size=30522, max_seq_len=77, 82 | device="cuda",use_tokenizer=True, embedding_dropout=0.0): 83 | super().__init__() 84 | self.use_tknz_fn = use_tokenizer 85 | if self.use_tknz_fn: 86 | self.tknz_fn = BERTTokenizer(vq_interface=False, max_length=max_seq_len) 87 | self.device = device 88 | self.transformer = TransformerWrapper(num_tokens=vocab_size, max_seq_len=max_seq_len, 89 | attn_layers=Encoder(dim=n_embed, depth=n_layer), 90 | emb_dropout=embedding_dropout) 91 | 92 | def forward(self, text): 93 | if self.use_tknz_fn: 94 | tokens = self.tknz_fn(text)#.to(self.device) 95 | else: 96 | tokens = text 97 | z = self.transformer(tokens, return_embeddings=True) 98 | return z 99 | 100 | def encode(self, text): 101 | # output of length 77 102 | return self(text) 103 | 104 | 105 | class SpatialRescaler(nn.Module): 106 | def __init__(self, 107 | n_stages=1, 108 | method='bilinear', 109 | multiplier=0.5, 110 | in_channels=3, 111 | out_channels=None, 112 | bias=False): 113 | super().__init__() 114 | self.n_stages = n_stages 115 | assert self.n_stages >= 0 116 | assert method in ['nearest','linear','bilinear','trilinear','bicubic','area'] 117 | self.multiplier = multiplier 118 | self.interpolator = partial(torch.nn.functional.interpolate, mode=method) 119 | self.remap_output = out_channels is not None 120 | if self.remap_output: 121 | print(f'Spatial Rescaler mapping from {in_channels} to {out_channels} channels after resizing.') 122 | self.channel_mapper = nn.Conv2d(in_channels,out_channels,1,bias=bias) 123 | 124 | def forward(self,x): 125 | for stage in range(self.n_stages): 126 | x = self.interpolator(x, scale_factor=self.multiplier) 127 | 128 | 129 | if self.remap_output: 130 | x = self.channel_mapper(x) 131 | return x 132 | 133 | def encode(self, x): 134 | return self(x) 135 | -------------------------------------------------------------------------------- /guided_diffusion/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Codebase for "Improved Denoising Diffusion Probabilistic Models". 3 | """ 4 | -------------------------------------------------------------------------------- /guided_diffusion/dist_util.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers for distributed training. 3 | """ 4 | 5 | import io 6 | import os 7 | import socket 8 | 9 | import blobfile as bf 10 | import torch as th 11 | import torch.distributed as dist 12 | from mpi4py import MPI 13 | 14 | # Change this to reflect your cluster layout. 15 | # The GPU for a given rank is (rank % GPUS_PER_NODE). 16 | GPUS_PER_NODE = MPI.COMM_WORLD.Get_size() 17 | 18 | SETUP_RETRY_COUNT = 3 19 | 20 | 21 | def setup_dist(): 22 | """ 23 | Setup a distributed process group. 24 | """ 25 | if dist.is_initialized(): 26 | return 27 | os.environ["CUDA_VISIBLE_DEVICES"] = f"{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}" 28 | 29 | comm = MPI.COMM_WORLD 30 | backend = "gloo" if not th.cuda.is_available() else "nccl" 31 | 32 | if backend == "gloo": 33 | hostname = "localhost" 34 | else: 35 | hostname = socket.gethostbyname(socket.getfqdn()) 36 | os.environ["MASTER_ADDR"] = comm.bcast(hostname, root=0) 37 | os.environ["RANK"] = str(comm.rank) 38 | os.environ["WORLD_SIZE"] = str(comm.size) 39 | 40 | port = comm.bcast(_find_free_port(), root=0) 41 | os.environ["MASTER_PORT"] = str(port) 42 | dist.init_process_group(backend=backend, init_method="env://") 43 | 44 | 45 | def dev(): 46 | """ 47 | Get the device to use for torch.distributed. 48 | """ 49 | if th.cuda.is_available(): 50 | return th.device(f"cuda") 51 | return th.device("cpu") 52 | 53 | 54 | def load_state_dict(path, **kwargs): 55 | with bf.BlobFile(path, "rb") as f: 56 | data = f.read() 57 | return th.load(io.BytesIO(data), **kwargs) 58 | 59 | 60 | def sync_params(params): 61 | """ 62 | Synchronize a sequence of Tensors across ranks from rank 0. 63 | """ 64 | for p in params: 65 | with th.no_grad(): 66 | dist.broadcast(p, 0) 67 | 68 | 69 | def _find_free_port(): 70 | try: 71 | s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 72 | s.bind(("", 0)) 73 | s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) 74 | return s.getsockname()[1] 75 | finally: 76 | s.close() 77 | -------------------------------------------------------------------------------- /guided_diffusion/fp16_util.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers to train with 16-bit precision. 3 | """ 4 | 5 | import numpy as np 6 | import torch as th 7 | import torch.nn as nn 8 | from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors 9 | 10 | from . import logger 11 | 12 | INITIAL_LOG_LOSS_SCALE = 20.0 13 | 14 | 15 | def convert_module_to_f16(l): 16 | """ 17 | Convert primitive modules to float16. 18 | """ 19 | if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): 20 | l.weight.data = l.weight.data.half() 21 | if l.bias is not None: 22 | l.bias.data = l.bias.data.half() 23 | 24 | 25 | def convert_module_to_f32(l): 26 | """ 27 | Convert primitive modules to float32, undoing convert_module_to_f16(). 28 | """ 29 | if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): 30 | l.weight.data = l.weight.data.float() 31 | if l.bias is not None: 32 | l.bias.data = l.bias.data.float() 33 | 34 | 35 | def make_master_params(param_groups_and_shapes): 36 | """ 37 | Copy model parameters into a (differently-shaped) list of full-precision 38 | parameters. 39 | """ 40 | master_params = [] 41 | for param_group, shape in param_groups_and_shapes: 42 | master_param = nn.Parameter( 43 | _flatten_dense_tensors( 44 | [param.detach().float() for (_, param) in param_group] 45 | ).view(shape) 46 | ) 47 | master_param.requires_grad = True 48 | master_params.append(master_param) 49 | return master_params 50 | 51 | 52 | def model_grads_to_master_grads(param_groups_and_shapes, master_params): 53 | """ 54 | Copy the gradients from the model parameters into the master parameters 55 | from make_master_params(). 56 | """ 57 | for master_param, (param_group, shape) in zip( 58 | master_params, param_groups_and_shapes 59 | ): 60 | master_param.grad = _flatten_dense_tensors( 61 | [param_grad_or_zeros(param) for (_, param) in param_group] 62 | ).view(shape) 63 | 64 | 65 | def master_params_to_model_params(param_groups_and_shapes, master_params): 66 | """ 67 | Copy the master parameter data back into the model parameters. 68 | """ 69 | # Without copying to a list, if a generator is passed, this will 70 | # silently not copy any parameters. 71 | for master_param, (param_group, _) in zip(master_params, param_groups_and_shapes): 72 | for (_, param), unflat_master_param in zip( 73 | param_group, unflatten_master_params(param_group, master_param.view(-1)) 74 | ): 75 | param.detach().copy_(unflat_master_param) 76 | 77 | 78 | def unflatten_master_params(param_group, master_param): 79 | return _unflatten_dense_tensors(master_param, [param for (_, param) in param_group]) 80 | 81 | 82 | def get_param_groups_and_shapes(named_model_params): 83 | named_model_params = list(named_model_params) 84 | scalar_vector_named_params = ( 85 | [(n, p) for (n, p) in named_model_params if p.ndim <= 1], 86 | (-1), 87 | ) 88 | matrix_named_params = ( 89 | [(n, p) for (n, p) in named_model_params if p.ndim > 1], 90 | (1, -1), 91 | ) 92 | return [scalar_vector_named_params, matrix_named_params] 93 | 94 | 95 | def master_params_to_state_dict( 96 | model, param_groups_and_shapes, master_params, use_fp16 97 | ): 98 | if use_fp16: 99 | state_dict = model.state_dict() 100 | for master_param, (param_group, _) in zip( 101 | master_params, param_groups_and_shapes 102 | ): 103 | for (name, _), unflat_master_param in zip( 104 | param_group, unflatten_master_params(param_group, master_param.view(-1)) 105 | ): 106 | assert name in state_dict 107 | state_dict[name] = unflat_master_param 108 | else: 109 | state_dict = model.state_dict() 110 | for i, (name, _value) in enumerate(model.named_parameters()): 111 | assert name in state_dict 112 | state_dict[name] = master_params[i] 113 | return state_dict 114 | 115 | 116 | def state_dict_to_master_params(model, state_dict, use_fp16): 117 | if use_fp16: 118 | named_model_params = [ 119 | (name, state_dict[name]) for name, _ in model.named_parameters() 120 | ] 121 | param_groups_and_shapes = get_param_groups_and_shapes(named_model_params) 122 | master_params = make_master_params(param_groups_and_shapes) 123 | else: 124 | master_params = [state_dict[name] for name, _ in model.named_parameters()] 125 | return master_params 126 | 127 | 128 | def zero_master_grads(master_params): 129 | for param in master_params: 130 | param.grad = None 131 | 132 | 133 | def zero_grad(model_params): 134 | for param in model_params: 135 | # Taken from https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.add_param_group 136 | if param.grad is not None: 137 | param.grad.detach_() 138 | param.grad.zero_() 139 | 140 | 141 | def param_grad_or_zeros(param): 142 | if param.grad is not None: 143 | return param.grad.data.detach() 144 | else: 145 | return th.zeros_like(param) 146 | 147 | 148 | class MixedPrecisionTrainer: 149 | def __init__( 150 | self, 151 | *, 152 | model, 153 | use_fp16=False, 154 | fp16_scale_growth=1e-3, 155 | initial_lg_loss_scale=INITIAL_LOG_LOSS_SCALE, 156 | ): 157 | self.model = model 158 | self.use_fp16 = use_fp16 159 | self.fp16_scale_growth = fp16_scale_growth 160 | 161 | self.model_params = list(self.model.parameters()) 162 | self.master_params = self.model_params 163 | self.param_groups_and_shapes = None 164 | self.lg_loss_scale = initial_lg_loss_scale 165 | 166 | if self.use_fp16: 167 | self.param_groups_and_shapes = get_param_groups_and_shapes( 168 | self.model.named_parameters() 169 | ) 170 | self.master_params = make_master_params(self.param_groups_and_shapes) 171 | self.model.convert_to_fp16() 172 | 173 | def zero_grad(self): 174 | zero_grad(self.model_params) 175 | 176 | def backward(self, loss: th.Tensor): 177 | if self.use_fp16: 178 | loss_scale = 2 ** self.lg_loss_scale 179 | (loss * loss_scale).backward() 180 | else: 181 | loss.backward() 182 | 183 | def optimize(self, opt: th.optim.Optimizer): 184 | if self.use_fp16: 185 | return self._optimize_fp16(opt) 186 | else: 187 | return self._optimize_normal(opt) 188 | 189 | def _optimize_fp16(self, opt: th.optim.Optimizer): 190 | logger.logkv_mean("lg_loss_scale", self.lg_loss_scale) 191 | model_grads_to_master_grads(self.param_groups_and_shapes, self.master_params) 192 | grad_norm, param_norm = self._compute_norms(grad_scale=2 ** self.lg_loss_scale) 193 | if check_overflow(grad_norm): 194 | self.lg_loss_scale -= 1 195 | logger.log(f"Found NaN, decreased lg_loss_scale to {self.lg_loss_scale}") 196 | zero_master_grads(self.master_params) 197 | return False 198 | 199 | logger.logkv_mean("grad_norm", grad_norm) 200 | logger.logkv_mean("param_norm", param_norm) 201 | 202 | self.master_params[0].grad.mul_(1.0 / (2 ** self.lg_loss_scale)) 203 | opt.step() 204 | zero_master_grads(self.master_params) 205 | master_params_to_model_params(self.param_groups_and_shapes, self.master_params) 206 | self.lg_loss_scale += self.fp16_scale_growth 207 | return True 208 | 209 | def _optimize_normal(self, opt: th.optim.Optimizer): 210 | grad_norm, param_norm = self._compute_norms() 211 | logger.logkv_mean("grad_norm", grad_norm) 212 | logger.logkv_mean("param_norm", param_norm) 213 | opt.step() 214 | return True 215 | 216 | def _compute_norms(self, grad_scale=1.0): 217 | grad_norm = 0.0 218 | param_norm = 0.0 219 | for p in self.master_params: 220 | with th.no_grad(): 221 | param_norm += th.norm(p, p=2, dtype=th.float32).item() ** 2 222 | if p.grad is not None: 223 | grad_norm += th.norm(p.grad, p=2, dtype=th.float32).item() ** 2 224 | return np.sqrt(grad_norm) / grad_scale, np.sqrt(param_norm) 225 | 226 | def master_params_to_state_dict(self, master_params): 227 | return master_params_to_state_dict( 228 | self.model, self.param_groups_and_shapes, master_params, self.use_fp16 229 | ) 230 | 231 | def state_dict_to_master_params(self, state_dict): 232 | return state_dict_to_master_params(self.model, state_dict, self.use_fp16) 233 | 234 | 235 | def check_overflow(value): 236 | return (value == float("inf")) or (value == -float("inf")) or (value != value) 237 | -------------------------------------------------------------------------------- /guided_diffusion/image_text_datasets.py: -------------------------------------------------------------------------------- 1 | import io 2 | import math 3 | import random 4 | from pathlib import Path 5 | 6 | import blobfile as bf 7 | import numpy as np 8 | import webdataset as wds 9 | from braceexpand import braceexpand 10 | from mpi4py import MPI 11 | from PIL import Image 12 | from torch.utils.data import DataLoader, Dataset 13 | 14 | 15 | def _list_image_files_recursively(data_dir): 16 | results = [] 17 | for entry in sorted(bf.listdir(data_dir)): 18 | full_path = bf.join(data_dir, entry) 19 | entry = entry.split(".") 20 | ext = entry[-1].strip() 21 | filename = entry[0] 22 | if ext and ext.lower() in ["jpg", "jpeg", "png", "gif", "webp"]: 23 | text_path = bf.join(data_dir, filename + ".txt") 24 | if bf.exists(text_path): 25 | results.append((full_path, text_path)) 26 | elif bf.isdir(full_path): 27 | results.extend(_list_image_files_recursively(full_path)) 28 | return results 29 | 30 | 31 | class CaptionedImageDataset(Dataset): 32 | def __init__( 33 | self, 34 | resolution, 35 | file_paths, 36 | shard=0, 37 | num_shards=1, 38 | random_crop=False, 39 | random_flip=True, 40 | ): 41 | super().__init__() 42 | self.resolution = resolution 43 | self.local_files = file_paths[shard:][::num_shards] 44 | self.random_crop = random_crop 45 | self.random_flip = random_flip 46 | 47 | def __len__(self): 48 | return len(self.local_files) 49 | 50 | def __getitem__(self, idx): 51 | path = self.local_files[idx] 52 | with bf.BlobFile(path[0], "rb") as f: 53 | pil_image = Image.open(f) 54 | pil_image.load() 55 | pil_image = pil_image.convert("RGB") 56 | 57 | if self.random_crop: 58 | arr = random_crop_arr(pil_image, self.resolution) 59 | else: 60 | arr = center_crop_arr(pil_image, self.resolution) 61 | 62 | if self.random_flip and random.random() < 0.5: 63 | arr = arr[:, ::-1] 64 | 65 | arr = arr.astype(np.float32) / 127.5 - 1 66 | 67 | with bf.BlobFile(path[1], "r") as f: 68 | text = f.read().strip() 69 | 70 | return np.transpose(arr, [2, 0, 1]), text 71 | 72 | 73 | def load_data( 74 | *, 75 | data_dir, 76 | batch_size, 77 | random_crop=False, 78 | random_flip=True, 79 | image_key="jpg", 80 | caption_key="txt", 81 | cache_dir=None, 82 | epochs=None, 83 | shard_size=None, 84 | ): 85 | """ 86 | For a dataset, create a generator over (images, kwargs) pairs. 87 | 88 | Each images is an NCHW float tensor, and the kwargs dict contains zero or 89 | more keys, each of which map to a batched Tensor of their own. 90 | The kwargs dict can be used for class labels, in which case the key is "y" 91 | and the values are integer tensors of class labels. 92 | 93 | :param data_dir: a dataset directory. 94 | :param batch_size: the batch size of each returned pair. 95 | :param deterministic: if True, yield results in a deterministic order. 96 | :param random_crop: if True, randomly crop the images for augmentation. 97 | :param random_flip: if True, randomly flip the images for augmentation. 98 | """ 99 | if not data_dir: 100 | raise ValueError("unspecified data directory") 101 | 102 | if ".tar" not in data_dir: 103 | print( 104 | f"Detected COCO-style (.txt/.jpg) dataset. Using CaptionImageLoader on {data_dir}." 105 | ) 106 | all_files = _list_image_files_recursively(data_dir) 107 | print(f"Found {len(all_files)} files") 108 | assert len(all_files) > 0, "no files found" 109 | dataset = CaptionedImageDataset( 110 | 256, 111 | all_files, 112 | shard=MPI.COMM_WORLD.Get_rank(), 113 | num_shards=MPI.COMM_WORLD.Get_size(), 114 | random_crop=random_crop, 115 | random_flip=random_flip, 116 | ) 117 | loader = DataLoader( 118 | dataset, batch_size=batch_size, shuffle=True, num_workers=8, drop_last=True 119 | ) 120 | while True: 121 | yield from loader 122 | else: 123 | print( 124 | "Detected webdataset (.tar files) dataset. Using WebDatasetLoader on {data_dir}." 125 | ) 126 | wds_uris = parse_data_dir(data_dir) 127 | assert len(wds_uris) > 0, "no files found" 128 | print(f"Found {len(wds_uris)} tar files of total {len(wds_uris)}") 129 | 130 | dataset = load_webdataset( 131 | 256, # TODO 132 | wds_uris, 133 | random_crop=random_crop, 134 | random_flip=random_flip, 135 | myimg=image_key, 136 | mycap=caption_key, 137 | cache_dir=cache_dir, 138 | ) 139 | if epochs and shard_size: 140 | total_size = epochs * shard_size * len(wds_uris) 141 | print(f"Number of samples to be trained: {total_size}") 142 | dataset = dataset.shuffle(total_size) 143 | dataset = dataset.batched(batch_size) 144 | loader = wds.WebLoader(dataset, batch_size=None, shuffle=False) 145 | while True: 146 | yield from loader 147 | 148 | 149 | def center_crop_arr(pil_image, image_size): 150 | # We are not on a new enough PIL to support the `reducing_gap` 151 | # argument, which uses BOX downsampling at powers of two first. 152 | # Thus, we do it by hand to improve downsample quality. 153 | while min(*pil_image.size) >= 2 * image_size: 154 | pil_image = pil_image.resize( 155 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX 156 | ) 157 | 158 | scale = image_size / min(*pil_image.size) 159 | pil_image = pil_image.resize( 160 | tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC 161 | ) 162 | 163 | arr = np.array(pil_image) 164 | crop_y = (arr.shape[0] - image_size) // 2 165 | crop_x = (arr.shape[1] - image_size) // 2 166 | return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] 167 | 168 | 169 | def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0): 170 | min_smaller_dim_size = math.ceil(image_size / max_crop_frac) 171 | max_smaller_dim_size = math.ceil(image_size / min_crop_frac) 172 | smaller_dim_size = random.randrange(min_smaller_dim_size, max_smaller_dim_size + 1) 173 | 174 | # We are not on a new enough PIL to support the `reducing_gap` 175 | # argument, which uses BOX downsampling at powers of two first. 176 | # Thus, we do it by hand to improve downsample quality. 177 | while min(*pil_image.size) >= 2 * smaller_dim_size: 178 | pil_image = pil_image.resize( 179 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX 180 | ) 181 | 182 | scale = smaller_dim_size / min(*pil_image.size) 183 | pil_image = pil_image.resize( 184 | tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC 185 | ) 186 | 187 | arr = np.array(pil_image) 188 | crop_y = random.randrange(arr.shape[0] - image_size + 1) 189 | crop_x = random.randrange(arr.shape[1] - image_size + 1) 190 | return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] 191 | 192 | 193 | def clean_caption(caption): 194 | caption = caption.decode("utf-8") 195 | caption = ( 196 | caption.replace("\n", " ") 197 | .replace("\t", " ") 198 | .replace("\r", " ") 199 | .replace(" ", " ") 200 | ) 201 | caption = caption.strip() 202 | return caption 203 | 204 | 205 | def load_webdataset( 206 | resolution, 207 | file_paths, 208 | random_crop=False, 209 | random_flip=False, 210 | myimg="jpg", 211 | mycap="txt", 212 | cache_dir=None, 213 | ): 214 | def bytes_to_pil_image(item): 215 | pil_image = Image.open(io.BytesIO(item)).convert("RGB") 216 | pil_image.load() 217 | return pil_image 218 | 219 | def filter_by_item(item): 220 | if mycap not in item: 221 | return False 222 | if myimg not in item: 223 | return False 224 | return True 225 | 226 | def pil_transform_to_np(arr): 227 | if random_crop: 228 | arr = random_crop_arr( 229 | arr, resolution, min_crop_frac=0.95 230 | ) # TODO make this a param 231 | else: 232 | arr = center_crop_arr(arr, resolution) 233 | if random_flip and random.random() < 0.5: 234 | arr = arr[:, ::-1] 235 | arr = arr.astype(np.float32) / 127.5 - 1 236 | return np.transpose(arr, [2, 0, 1]) 237 | 238 | image_text_mapping = {myimg: bytes_to_pil_image, mycap: clean_caption} 239 | image_mapping = {myimg: pil_transform_to_np} 240 | dataset = wds.WebDataset( 241 | urls=file_paths, 242 | handler=wds.warn_and_continue, 243 | cache_dir=cache_dir, 244 | # shardshuffle=True, 245 | # nodesplitter=wds.split_by_worker, 246 | ) 247 | filtered_dataset = dataset.select(filter_by_item) 248 | dataset = ( 249 | filtered_dataset.map_dict(**image_text_mapping) 250 | .map_dict(**image_mapping) 251 | .to_tuple(myimg, mycap) 252 | ) 253 | return dataset 254 | 255 | 256 | def parse_data_dir(data_dir): 257 | if Path(data_dir).is_dir(): 258 | wds_uris = [ 259 | str(p) for p in Path(data_dir).glob("**/*") if ".tar" in str(p).lower() 260 | ] 261 | assert ( 262 | len(wds_uris) > 0 263 | ), "The directory ({}) does not contain any WebDataset/.tar files.".format( 264 | data_dir 265 | ) 266 | print( 267 | "Found {} WebDataset .tar(.gz) file(s) under given path {}!".format( 268 | len(wds_uris), data_dir 269 | ) 270 | ) 271 | elif "s3://" in data_dir.lower(): 272 | data_dir = f"pipe:aws s3 cp {data_dir} -" 273 | elif ("http://" in data_dir.lower()) | ("https://" in data_dir.lower()): 274 | wds_uris = f"pipe:curl -L -s {data_dir} || true" 275 | print("Found {} http(s) link under given path!".format(len(wds_uris), data_dir)) 276 | elif "gs://" in data_dir.lower(): 277 | wds_uris = f"pipe:gsutil cat {data_dir} || true" 278 | print("Found {} GCS link under given path!".format(len(wds_uris), data_dir)) 279 | 280 | if ".tar" in data_dir: 281 | wds_uris = braceexpand(data_dir) 282 | print("Found WebDataset .tar(.gz) file under given path {}!".format(data_dir)) 283 | return wds_uris 284 | -------------------------------------------------------------------------------- /guided_diffusion/inpaint_util.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import re 4 | import typing 5 | import unicodedata 6 | from pathlib import Path 7 | 8 | import torch 9 | from dist.clip_custom import clip 10 | from PIL import Image 11 | from torchvision import transforms 12 | from torchvision.transforms import functional as TF 13 | 14 | from guided_diffusion.predict_util import ( 15 | average_prompt_embed_with_aesthetic_embed, 16 | bert_encode_cfg, 17 | load_aesthetic_vit_l_14_embed, 18 | load_bert, 19 | load_clip_model_and_transform, 20 | load_diffusion_model, 21 | load_vae, 22 | pack_model_kwargs, 23 | prepare_edit, 24 | ) 25 | from guided_diffusion.script_util import create_gaussian_diffusion 26 | 27 | 28 | def set_requires_grad(model, value): 29 | for param in model.parameters(): 30 | param.requires_grad = value 31 | 32 | 33 | normalize = transforms.Normalize( 34 | mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711] 35 | ) 36 | 37 | os.environ[ 38 | "TOKENIZERS_PARALLELISM" 39 | ] = "false" # required to avoid errors with transformers lib 40 | 41 | 42 | KL_PATH = "kl-f8.pt" 43 | BERT_PATH = "bert.pt" 44 | 45 | 46 | def prepare_inpaint_models( 47 | inpaint_model_path: str = "inpaint.pt", device: str = "cuda", use_fp16: bool = False 48 | ): 49 | device = torch.device(device) 50 | print(f"Loading latent diffusion model from {inpaint_model_path}") 51 | inpaint_model, inpaint_model_config, inpaint_diffusion = load_diffusion_model( 52 | model_path=inpaint_model_path, 53 | steps="1000", # Init method requires steps, although we can modify it during inference as well. 54 | use_fp16=use_fp16, 55 | device=device, 56 | ) 57 | 58 | print(f"Loading VAE from {KL_PATH}") 59 | vae_backbone = load_vae(kl_path=KL_PATH, device=device, use_fp16=use_fp16) 60 | 61 | print(f"Loading CLIP") 62 | clip_model, clip_preprocess = load_clip_model_and_transform(device) 63 | 64 | print(f"Loading BERT text encoder from {BERT_PATH}") 65 | bert = load_bert(BERT_PATH, device, use_fp16=use_fp16) 66 | return dict( 67 | inpaint_model=inpaint_model, 68 | inpaint_model_config=inpaint_model_config, 69 | inpaint_diffusion=inpaint_diffusion, 70 | vae_backbone=vae_backbone, 71 | clip_model=clip_model, 72 | clip_preprocess=clip_preprocess, 73 | bert=bert, 74 | ) 75 | 76 | 77 | def slugify(value, allow_unicode=False): 78 | """ 79 | Taken from https://github.com/django/django/blob/master/django/utils/text.py 80 | Convert to ASCII if 'allow_unicode' is False. Convert spaces or repeated 81 | dashes to single dashes. Remove characters that aren't alphanumerics, 82 | underscores, or hyphens. Convert to lowercase. Also strip leading and 83 | trailing whitespace, dashes, and underscores. 84 | """ 85 | value = str(value) 86 | if allow_unicode: 87 | value = unicodedata.normalize("NFKC", value) 88 | else: 89 | value = ( 90 | unicodedata.normalize("NFKD", value) 91 | .encode("ascii", "ignore") 92 | .decode("ascii") 93 | ) 94 | value = re.sub(r"[^\w\s-]", "", value.lower()) 95 | return re.sub(r"[-\s]+", "-", value).strip("-_") 96 | 97 | 98 | def sample_inpaint( 99 | prompt: str, 100 | negative: str = "", 101 | init_image: str = None, 102 | mask: str = None, 103 | steps: int = 100, 104 | init_skip_fraction: float = 0.5, 105 | width: int = 256, 106 | height: int = 256, 107 | batch_size: int = 1, 108 | intermediate_outputs: bool = False, 109 | guidance_scale: float = 0.0, 110 | aesthetic_rating: int = 9, 111 | aesthetic_weight: float = 0.0, 112 | device: str = "cuda", 113 | use_fp16: bool = False, 114 | seed: int = 0, 115 | output_dir: str = "inpaint_outputs", 116 | loaded_models: typing.Dict = None, 117 | ): 118 | """Predict a normal distribution from a prompt. 119 | 120 | Args: 121 | prompt: The prompt to use. 122 | negative: The negative prompt to use. 123 | init_image: The image to use as the initial image. 124 | steps: The number of steps to run the model. 125 | mask: The mask to use for the initial image. 126 | init_skip_fraction: The fraction of timesteps to skip when using init_image. 127 | width: The width of the output image. 128 | height: The height of the output image. 129 | batch_size: The batch size to use. 130 | intermediate_outputs: Whether to save intermediate outputs. 131 | guidance_scale: The scale to use for guidance. 132 | aesthetic_rating: The rating to use for the aesthetic embedding. 133 | aesthetic_weight: The weight to use for the aesthetic embedding. 134 | device: The device to use. 135 | use_fp16: Whether to use fp16. 136 | loaded_models: A dictionary of models pre-loaded to `device` to use for inference. Keys are: 137 | - inpaint_model 138 | - inpaint_model_config 139 | - inpaint_diffusion 140 | - vae_backbone 141 | - clip_model 142 | - clip_preprocess 143 | - bert 144 | 145 | Returns: 146 | A generator that yields a list of paths to images. 147 | """ 148 | prompt_dir = Path(output_dir).joinpath(slugify(prompt)) 149 | prompt_dir.mkdir(parents=True, exist_ok=True) 150 | if seed > 0: 151 | torch.manual_seed(seed) 152 | else: 153 | seed = random.randint(0, 2**32) 154 | torch.manual_seed(seed) 155 | print(f"Using seed {seed}") 156 | 157 | if loaded_models is None: 158 | loaded_models = prepare_inpaint_models(device=device, use_fp16=use_fp16) 159 | else: 160 | print("Using preloaded models") 161 | 162 | model_config = loaded_models["inpaint_model_config"] 163 | clip_model = loaded_models["clip_model"] 164 | bert = loaded_models["bert"] 165 | vq_decoder = loaded_models["vae_backbone"] 166 | inpaint_model = loaded_models["inpaint_model"] 167 | 168 | # Create diffusion manually so we don't re-init the model just to change timestep_respacing 169 | model_config["timestep_respacing"] = str(steps) 170 | diffusion = create_gaussian_diffusion( 171 | steps=model_config["diffusion_steps"], 172 | learn_sigma=model_config["learn_sigma"], 173 | noise_schedule=model_config["noise_schedule"], 174 | use_kl=model_config["use_kl"], 175 | predict_xstart=model_config["predict_xstart"], 176 | rescale_timesteps=model_config["rescale_timesteps"], 177 | timestep_respacing=model_config["timestep_respacing"], 178 | ) 179 | 180 | # Text Setup 181 | print(f"Encoding text embeddings with {prompt} dimensions") 182 | text_emb, text_blank = bert_encode_cfg(prompt, negative, batch_size, device, bert) 183 | 184 | text_tokens = clip.tokenize([prompt] * batch_size, truncate=True).to(device) 185 | negative_tokens = clip.tokenize([negative] * batch_size, truncate=True).to(device) 186 | text_emb_clip = clip_model.encode_text(text_tokens).to(device).float() 187 | text_emb_clip_blank = clip_model.encode_text(negative_tokens).to(device).float() 188 | text_emb_norm = text_emb_clip[0] / text_emb_clip[0].norm(dim=-1, keepdim=True) 189 | print( 190 | f"Using aesthetic embedding {aesthetic_rating} with weight {aesthetic_weight}" 191 | ) 192 | text_emb_clip_aesthetic = load_aesthetic_vit_l_14_embed(rating=aesthetic_rating).to( 193 | device 194 | ) 195 | text_emb_clip = average_prompt_embed_with_aesthetic_embed( 196 | text_emb_clip, text_emb_clip_aesthetic, aesthetic_weight 197 | ) 198 | 199 | # Image Setup 200 | 201 | init = None 202 | init_skip_fraction = 0.0 203 | init_skip_timesteps = 0 204 | 205 | image_embed = torch.zeros(batch_size * 2, 4, height // 8, width // 8, device=device) 206 | if init_image and mask: # if both are provided, the user is inpainting. 207 | print(f"Using inpaint model with image: {init_image}") 208 | image_embed = prepare_edit( 209 | vq_decoder, str(init_image), width, height, device=device 210 | ) 211 | mask_image = Image.open(mask).convert("L") 212 | mask_image = mask_image.resize((width // 8, height // 8), Image.ANTIALIAS) 213 | mask = transforms.ToTensor()(mask_image).unsqueeze(0).to(device) 214 | mask1 = mask > 0.5 215 | mask1 = mask1.float() 216 | image_embed *= mask1 217 | image_embed = torch.cat(batch_size * 2 * [image_embed], dim=0) 218 | elif ( 219 | init_image 220 | ): # if just the image is provided, the user wants to use the image as the init image. 221 | if init_skip_fraction == 0.0: 222 | print(f"Must specify init_skip_fraction > 0.0 when using init_image.") 223 | print(f"Overriding init_skip_fraction to 0.5") 224 | init_skip_fraction = 0.5 225 | print( 226 | f"Loading initial image {init_image} with init_skip_fraction: {init_skip_fraction}" 227 | ) 228 | init = Image.open(init_image).convert("RGB") 229 | init = init.resize((int(width), int(height)), Image.LANCZOS) 230 | init = TF.to_tensor(init).to(device).unsqueeze(0).clamp(0, 1) 231 | if use_fp16: 232 | init = init.half() 233 | h = vq_decoder.encode(init * 2 - 1).sample() * 0.18215 234 | init = torch.cat(batch_size * 2 * [h], dim=0) 235 | # str to int * float -> float 236 | init_skip_timesteps = ( 237 | int(model_config["timestep_respacing"]) * init_skip_fraction 238 | ) 239 | # float to int 240 | init_skip_timesteps = int(init_skip_timesteps) 241 | 242 | # Prepare inputs 243 | kwargs = pack_model_kwargs( 244 | text_emb=text_emb, 245 | text_blank=text_blank, 246 | text_emb_clip=text_emb_clip, 247 | text_emb_clip_blank=text_emb_clip_blank, 248 | image_embed=image_embed, 249 | model_params=model_config, 250 | ) 251 | 252 | # Create a classifier-free guidance sampling function. 253 | @torch.cuda.amp.autocast(enabled=use_fp16) 254 | def model_fn(x_t, ts, **kwargs): 255 | half = x_t[: len(x_t) // 2] 256 | combined = torch.cat([half, half], dim=0) 257 | model_out = inpaint_model(combined, ts, **kwargs) 258 | eps, rest = model_out[:, :3], model_out[:, 3:] 259 | cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0) 260 | half_eps = uncond_eps + guidance_scale * (cond_eps - uncond_eps) 261 | eps = torch.cat([half_eps, half_eps], dim=0) 262 | return torch.cat([eps, rest], dim=1) 263 | 264 | @torch.cuda.amp.autocast(enabled=use_fp16) 265 | def save_sample(sample: torch.Tensor) -> typing.List[torch.Tensor]: 266 | """Save a sample of the model's output.""" 267 | final_outputs = [] 268 | for image in sample["pred_xstart"][:batch_size]: 269 | image /= 0.18215 270 | im = image.unsqueeze(0) 271 | out = vq_decoder.decode(im) 272 | final_outputs.append(out.squeeze(0).add(1).div(2).clamp(0, 1)) 273 | return final_outputs 274 | 275 | sample_fn = diffusion.plms_sample_loop_progressive 276 | samples = sample_fn( 277 | model_fn, 278 | (batch_size * 2, 4, int(height / 8), int(width / 8)), 279 | clip_denoised=False, 280 | model_kwargs=kwargs, 281 | cond_fn=None, 282 | device=device, 283 | progress=True, 284 | init_image=init, 285 | skip_timesteps=init_skip_timesteps, 286 | ) 287 | 288 | log_interval = 10 289 | print("Running diffusion...") 290 | for timestep_idx, sample in enumerate(samples): 291 | if ( 292 | timestep_idx % log_interval == 0 293 | and timestep_idx < diffusion.num_timesteps - 1 294 | and intermediate_outputs 295 | ): 296 | print(f"Timestep {timestep_idx+1} - saving sample/s") 297 | current_batch = save_sample(sample) 298 | current_batch_paths = [] 299 | for batch_idx, current_image in enumerate(current_batch): 300 | current_image_path = prompt_dir.joinpath( 301 | f"ts_{timestep_idx}-batch_{batch_idx}.png" 302 | ) 303 | current_batch_paths.append(current_image_path) 304 | TF.to_pil_image(current_image).save(current_image_path, optimize=True) 305 | yield current_batch_paths # List[str] 306 | 307 | print(f"Saving final sample/s") 308 | current_batch = save_sample(sample) 309 | current_batch_paths = [] 310 | for batch_idx, current_image in enumerate(current_batch): 311 | current_image_path = prompt_dir.joinpath( 312 | f"ts_{timestep_idx}-batch_{batch_idx}.png" 313 | ) 314 | current_batch_paths.append(current_image_path) 315 | TF.to_pil_image(current_image).save(current_image_path, optimize=True) 316 | yield current_batch_paths 317 | -------------------------------------------------------------------------------- /guided_diffusion/losses.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers for various likelihood-based losses. These are ported from the original 3 | Ho et al. diffusion models codebase: 4 | https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/utils.py 5 | """ 6 | 7 | import numpy as np 8 | import torch as th 9 | 10 | 11 | def normal_kl(mean1, logvar1, mean2, logvar2): 12 | """ 13 | Compute the KL divergence between two gaussians. 14 | 15 | Shapes are automatically broadcasted, so batches can be compared to 16 | scalars, among other use cases. 17 | """ 18 | tensor = None 19 | for obj in (mean1, logvar1, mean2, logvar2): 20 | if isinstance(obj, th.Tensor): 21 | tensor = obj 22 | break 23 | assert tensor is not None, "at least one argument must be a Tensor" 24 | 25 | # Force variances to be Tensors. Broadcasting helps convert scalars to 26 | # Tensors, but it does not work for th.exp(). 27 | logvar1, logvar2 = [ 28 | x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor) 29 | for x in (logvar1, logvar2) 30 | ] 31 | 32 | return 0.5 * ( 33 | -1.0 34 | + logvar2 35 | - logvar1 36 | + th.exp(logvar1 - logvar2) 37 | + ((mean1 - mean2) ** 2) * th.exp(-logvar2) 38 | ) 39 | 40 | 41 | def approx_standard_normal_cdf(x): 42 | """ 43 | A fast approximation of the cumulative distribution function of the 44 | standard normal. 45 | """ 46 | return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3)))) 47 | 48 | 49 | def discretized_gaussian_log_likelihood(x, *, means, log_scales): 50 | """ 51 | Compute the log-likelihood of a Gaussian distribution discretizing to a 52 | given image. 53 | 54 | :param x: the target images. It is assumed that this was uint8 values, 55 | rescaled to the range [-1, 1]. 56 | :param means: the Gaussian mean Tensor. 57 | :param log_scales: the Gaussian log stddev Tensor. 58 | :return: a tensor like x of log probabilities (in nats). 59 | """ 60 | assert x.shape == means.shape == log_scales.shape 61 | centered_x = x - means 62 | inv_stdv = th.exp(-log_scales) 63 | plus_in = inv_stdv * (centered_x + 1.0 / 255.0) 64 | cdf_plus = approx_standard_normal_cdf(plus_in) 65 | min_in = inv_stdv * (centered_x - 1.0 / 255.0) 66 | cdf_min = approx_standard_normal_cdf(min_in) 67 | log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12)) 68 | log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12)) 69 | cdf_delta = cdf_plus - cdf_min 70 | log_probs = th.where( 71 | x < -0.999, 72 | log_cdf_plus, 73 | th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))), 74 | ) 75 | assert log_probs.shape == x.shape 76 | return log_probs 77 | -------------------------------------------------------------------------------- /guided_diffusion/nn.py: -------------------------------------------------------------------------------- 1 | """ 2 | Various utilities for neural networks. 3 | """ 4 | 5 | import math 6 | 7 | import torch as th 8 | import torch.nn as nn 9 | import torch.nn.functional as F 10 | 11 | 12 | class GroupNorm32(nn.GroupNorm): 13 | def __init__(self, num_groups, num_channels, swish, eps=1e-5): 14 | super().__init__(num_groups=num_groups, num_channels=num_channels, eps=eps) 15 | self.swish = swish 16 | 17 | def forward(self, x): 18 | y = super().forward(x.float()).to(x.dtype) 19 | if self.swish == 1.0: 20 | y = F.silu(y) 21 | elif self.swish: 22 | y = y * F.sigmoid(y * float(self.swish)) 23 | return y 24 | 25 | 26 | def conv_nd(dims, *args, **kwargs): 27 | """ 28 | Create a 1D, 2D, or 3D convolution module. 29 | """ 30 | if dims == 1: 31 | return nn.Conv1d(*args, **kwargs) 32 | elif dims == 2: 33 | return nn.Conv2d(*args, **kwargs) 34 | elif dims == 3: 35 | return nn.Conv3d(*args, **kwargs) 36 | raise ValueError(f"unsupported dimensions: {dims}") 37 | 38 | 39 | def linear(*args, **kwargs): 40 | """ 41 | Create a linear module. 42 | """ 43 | return nn.Linear(*args, **kwargs) 44 | 45 | 46 | def avg_pool_nd(dims, *args, **kwargs): 47 | """ 48 | Create a 1D, 2D, or 3D average pooling module. 49 | """ 50 | if dims == 1: 51 | return nn.AvgPool1d(*args, **kwargs) 52 | elif dims == 2: 53 | return nn.AvgPool2d(*args, **kwargs) 54 | elif dims == 3: 55 | return nn.AvgPool3d(*args, **kwargs) 56 | raise ValueError(f"unsupported dimensions: {dims}") 57 | 58 | 59 | def update_ema(target_params, source_params, rate=0.99): 60 | """ 61 | Update target parameters to be closer to those of source parameters using 62 | an exponential moving average. 63 | 64 | :param target_params: the target parameter sequence. 65 | :param source_params: the source parameter sequence. 66 | :param rate: the EMA rate (closer to 1 means slower). 67 | """ 68 | for targ, src in zip(target_params, source_params): 69 | targ.detach().mul_(rate).add_(src, alpha=1 - rate) 70 | 71 | 72 | def zero_module(module): 73 | """ 74 | Zero out the parameters of a module and return it. 75 | """ 76 | for p in module.parameters(): 77 | p.detach().zero_() 78 | return module 79 | 80 | 81 | def scale_module(module, scale): 82 | """ 83 | Scale the parameters of a module and return it. 84 | """ 85 | for p in module.parameters(): 86 | p.detach().mul_(scale) 87 | return module 88 | 89 | 90 | def mean_flat(tensor): 91 | """ 92 | Take the mean over all non-batch dimensions. 93 | """ 94 | return tensor.mean(dim=list(range(1, len(tensor.shape)))) 95 | 96 | 97 | def normalization(channels, swish=0.0): 98 | """ 99 | Make a standard normalization layer, with an optional swish activation. 100 | 101 | :param channels: number of input channels. 102 | :return: an nn.Module for normalization. 103 | """ 104 | return GroupNorm32(num_channels=channels, num_groups=32, swish=swish) 105 | 106 | 107 | # def timestep_embedding(timesteps, dim, max_period=10000): 108 | # """ 109 | # Create sinusoidal timestep embeddings. 110 | 111 | # :param timesteps: a 1-D Tensor of N indices, one per batch element. 112 | # These may be fractional. 113 | # :param dim: the dimension of the output. 114 | # :param max_period: controls the minimum frequency of the embeddings. 115 | # :return: an [N x dim] Tensor of positional embeddings. 116 | # """ 117 | # half = dim // 2 118 | # freqs = th.exp( 119 | # -math.log(max_period) * th.arange(start=0, end=half, dtype=th.float32) / half 120 | # ).to(device=timesteps.device) 121 | # args = timesteps[:, None].float() * freqs[None] 122 | # embedding = th.cat([th.cos(args), th.sin(args)], dim=-1) 123 | # if dim % 2: 124 | # embedding = th.cat([embedding, th.zeros_like(embedding[:, :1])], dim=-1) 125 | # return embedding 126 | 127 | 128 | def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False): 129 | """ 130 | Create sinusoidal timestep embeddings. 131 | :param timesteps: a 1-D Tensor of N indices, one per batch element. 132 | These may be fractional. 133 | :param dim: the dimension of the output. 134 | :param max_period: controls the minimum frequency of the embeddings. 135 | :return: an [N x dim] Tensor of positional embeddings. 136 | """ 137 | if not repeat_only: 138 | half = dim // 2 139 | freqs = th.exp( 140 | -math.log(max_period) 141 | * th.arange(start=0, end=half, dtype=th.float32) 142 | / half 143 | ).to(device=timesteps.device) 144 | args = timesteps[:, None].float() * freqs[None] 145 | embedding = th.cat([th.cos(args), th.sin(args)], dim=-1) 146 | if dim % 2: 147 | embedding = th.cat([embedding, th.zeros_like(embedding[:, :1])], dim=-1) 148 | else: 149 | embedding = repeat(timesteps, "b -> b d", d=dim) 150 | return embedding 151 | 152 | 153 | def checkpoint(func, inputs, params, flag): 154 | """ 155 | Evaluate a function without caching intermediate activations, allowing for 156 | reduced memory at the expense of extra compute in the backward pass. 157 | 158 | :param func: the function to evaluate. 159 | :param inputs: the argument sequence to pass to `func`. 160 | :param params: a sequence of parameters `func` depends on but does not 161 | explicitly take as arguments. 162 | :param flag: if False, disable gradient checkpointing. 163 | """ 164 | if flag: 165 | args = tuple(inputs) + tuple(params) 166 | return CheckpointFunction.apply(func, len(inputs), *args) 167 | else: 168 | return func(*inputs) 169 | 170 | 171 | class CheckpointFunction(th.autograd.Function): 172 | @staticmethod 173 | def forward(ctx, run_function, length, *args): 174 | ctx.run_function = run_function 175 | ctx.input_tensors = list(args[:length]) 176 | ctx.input_params = list(args[length:]) 177 | with th.no_grad(): 178 | output_tensors = ctx.run_function(*ctx.input_tensors) 179 | return output_tensors 180 | 181 | @staticmethod 182 | def backward(ctx, *output_grads): 183 | ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors] 184 | with th.enable_grad(): 185 | # Fixes a bug where the first op in run_function modifies the 186 | # Tensor storage in place, which is not allowed for detach()'d 187 | # Tensors. 188 | shallow_copies = [x.view_as(x) for x in ctx.input_tensors] 189 | output_tensors = ctx.run_function(*shallow_copies) 190 | input_grads = th.autograd.grad( 191 | output_tensors, 192 | ctx.input_tensors + ctx.input_params, 193 | output_grads, 194 | allow_unused=True, 195 | ) 196 | del ctx.input_tensors 197 | del ctx.input_params 198 | del output_tensors 199 | return (None, None) + input_grads 200 | -------------------------------------------------------------------------------- /guided_diffusion/resample.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | 3 | import numpy as np 4 | import torch as th 5 | import torch.distributed as dist 6 | 7 | 8 | def create_named_schedule_sampler(name, diffusion): 9 | """ 10 | Create a ScheduleSampler from a library of pre-defined samplers. 11 | 12 | :param name: the name of the sampler. 13 | :param diffusion: the diffusion object to sample for. 14 | """ 15 | if name == "uniform": 16 | return UniformSampler(diffusion) 17 | elif name == "loss-second-moment": 18 | return LossSecondMomentResampler(diffusion) 19 | else: 20 | raise NotImplementedError(f"unknown schedule sampler: {name}") 21 | 22 | 23 | class ScheduleSampler(ABC): 24 | """ 25 | A distribution over timesteps in the diffusion process, intended to reduce 26 | variance of the objective. 27 | 28 | By default, samplers perform unbiased importance sampling, in which the 29 | objective's mean is unchanged. 30 | However, subclasses may override sample() to change how the resampled 31 | terms are reweighted, allowing for actual changes in the objective. 32 | """ 33 | 34 | @abstractmethod 35 | def weights(self): 36 | """ 37 | Get a numpy array of weights, one per diffusion step. 38 | 39 | The weights needn't be normalized, but must be positive. 40 | """ 41 | 42 | def sample(self, batch_size, device): 43 | """ 44 | Importance-sample timesteps for a batch. 45 | 46 | :param batch_size: the number of timesteps. 47 | :param device: the torch device to save to. 48 | :return: a tuple (timesteps, weights): 49 | - timesteps: a tensor of timestep indices. 50 | - weights: a tensor of weights to scale the resulting losses. 51 | """ 52 | w = self.weights() 53 | p = w / np.sum(w) 54 | indices_np = np.random.choice(len(p), size=(batch_size,), p=p) 55 | indices = th.from_numpy(indices_np).long().to(device) 56 | weights_np = 1 / (len(p) * p[indices_np]) 57 | weights = th.from_numpy(weights_np).float().to(device) 58 | return indices, weights 59 | 60 | 61 | class UniformSampler(ScheduleSampler): 62 | def __init__(self, diffusion): 63 | self.diffusion = diffusion 64 | self._weights = np.ones([diffusion.num_timesteps]) 65 | 66 | def weights(self): 67 | return self._weights 68 | 69 | 70 | class LossAwareSampler(ScheduleSampler): 71 | def update_with_local_losses(self, local_ts, local_losses): 72 | """ 73 | Update the reweighting using losses from a model. 74 | 75 | Call this method from each rank with a batch of timesteps and the 76 | corresponding losses for each of those timesteps. 77 | This method will perform synchronization to make sure all of the ranks 78 | maintain the exact same reweighting. 79 | 80 | :param local_ts: an integer Tensor of timesteps. 81 | :param local_losses: a 1D Tensor of losses. 82 | """ 83 | batch_sizes = [ 84 | th.tensor([0], dtype=th.int32, device=local_ts.device) 85 | for _ in range(dist.get_world_size()) 86 | ] 87 | dist.all_gather( 88 | batch_sizes, 89 | th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device), 90 | ) 91 | 92 | # Pad all_gather batches to be the maximum batch size. 93 | batch_sizes = [x.item() for x in batch_sizes] 94 | max_bs = max(batch_sizes) 95 | 96 | timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes] 97 | loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes] 98 | dist.all_gather(timestep_batches, local_ts) 99 | dist.all_gather(loss_batches, local_losses) 100 | timesteps = [ 101 | x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs] 102 | ] 103 | losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]] 104 | self.update_with_all_losses(timesteps, losses) 105 | 106 | @abstractmethod 107 | def update_with_all_losses(self, ts, losses): 108 | """ 109 | Update the reweighting using losses from a model. 110 | 111 | Sub-classes should override this method to update the reweighting 112 | using losses from the model. 113 | 114 | This method directly updates the reweighting without synchronizing 115 | between workers. It is called by update_with_local_losses from all 116 | ranks with identical arguments. Thus, it should have deterministic 117 | behavior to maintain state across workers. 118 | 119 | :param ts: a list of int timesteps. 120 | :param losses: a list of float losses, one per timestep. 121 | """ 122 | 123 | 124 | class LossSecondMomentResampler(LossAwareSampler): 125 | def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001): 126 | self.diffusion = diffusion 127 | self.history_per_term = history_per_term 128 | self.uniform_prob = uniform_prob 129 | self._loss_history = np.zeros( 130 | [diffusion.num_timesteps, history_per_term], dtype=np.float64 131 | ) 132 | self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int) 133 | 134 | def weights(self): 135 | if not self._warmed_up(): 136 | return np.ones([self.diffusion.num_timesteps], dtype=np.float64) 137 | weights = np.sqrt(np.mean(self._loss_history**2, axis=-1)) 138 | weights /= np.sum(weights) 139 | weights *= 1 - self.uniform_prob 140 | weights += self.uniform_prob / len(weights) 141 | return weights 142 | 143 | def update_with_all_losses(self, ts, losses): 144 | for t, loss in zip(ts, losses): 145 | if self._loss_counts[t] == self.history_per_term: 146 | # Shift out the oldest loss term. 147 | self._loss_history[t, :-1] = self._loss_history[t, 1:] 148 | self._loss_history[t, -1] = loss 149 | else: 150 | self._loss_history[t, self._loss_counts[t]] = loss 151 | self._loss_counts[t] += 1 152 | 153 | def _warmed_up(self): 154 | return (self._loss_counts == self.history_per_term).all() 155 | -------------------------------------------------------------------------------- /guided_diffusion/respace.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch as th 3 | 4 | from .gaussian_diffusion import GaussianDiffusion 5 | 6 | 7 | def space_timesteps(num_timesteps, section_counts): 8 | """ 9 | Create a list of timesteps to use from an original diffusion process, 10 | given the number of timesteps we want to take from equally-sized portions 11 | of the original process. 12 | 13 | For example, if there's 300 timesteps and the section counts are [10,15,20] 14 | then the first 100 timesteps are strided to be 10 timesteps, the second 100 15 | are strided to be 15 timesteps, and the final 100 are strided to be 20. 16 | 17 | If the stride is a string starting with "ddim", then the fixed striding 18 | from the DDIM paper is used, and only one section is allowed. 19 | 20 | :param num_timesteps: the number of diffusion steps in the original 21 | process to divide up. 22 | :param section_counts: either a list of numbers, or a string containing 23 | comma-separated numbers, indicating the step count 24 | per section. As a special case, use "ddimN" where N 25 | is a number of steps to use the striding from the 26 | DDIM paper. 27 | :return: a set of diffusion steps from the original process to use. 28 | """ 29 | if isinstance(section_counts, str): 30 | if section_counts.startswith("ddim"): 31 | desired_count = int(section_counts[len("ddim") :]) 32 | for i in range(1, num_timesteps): 33 | if len(range(0, num_timesteps, i)) == desired_count: 34 | return set(range(0, num_timesteps, i)) 35 | raise ValueError( 36 | f"cannot create exactly {num_timesteps} steps with an integer stride" 37 | ) 38 | section_counts = [int(x) for x in section_counts.split(",")] 39 | size_per = num_timesteps // len(section_counts) 40 | extra = num_timesteps % len(section_counts) 41 | start_idx = 0 42 | all_steps = [] 43 | for i, section_count in enumerate(section_counts): 44 | size = size_per + (1 if i < extra else 0) 45 | if size < section_count: 46 | raise ValueError( 47 | f"cannot divide section of {size} steps into {section_count}" 48 | ) 49 | if section_count <= 1: 50 | frac_stride = 1 51 | else: 52 | frac_stride = (size - 1) / (section_count - 1) 53 | cur_idx = 0.0 54 | taken_steps = [] 55 | for _ in range(section_count): 56 | taken_steps.append(start_idx + round(cur_idx)) 57 | cur_idx += frac_stride 58 | all_steps += taken_steps 59 | start_idx += size 60 | return set(all_steps) 61 | 62 | 63 | class SpacedDiffusion(GaussianDiffusion): 64 | """ 65 | A diffusion process which can skip steps in a base diffusion process. 66 | 67 | :param use_timesteps: a collection (sequence or set) of timesteps from the 68 | original diffusion process to retain. 69 | :param kwargs: the kwargs to create the base diffusion process. 70 | """ 71 | 72 | def __init__(self, use_timesteps, **kwargs): 73 | self.use_timesteps = set(use_timesteps) 74 | self.timestep_map = [] 75 | self.original_num_steps = len(kwargs["betas"]) 76 | 77 | base_diffusion = GaussianDiffusion(**kwargs) # pylint: disable=missing-kwoa 78 | last_alpha_cumprod = 1.0 79 | new_betas = [] 80 | for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod): 81 | if i in self.use_timesteps: 82 | new_betas.append(1 - alpha_cumprod / last_alpha_cumprod) 83 | last_alpha_cumprod = alpha_cumprod 84 | self.timestep_map.append(i) 85 | kwargs["betas"] = np.array(new_betas) 86 | super().__init__(**kwargs) 87 | 88 | def p_mean_variance( 89 | self, model, *args, **kwargs 90 | ): # pylint: disable=signature-differs 91 | return super().p_mean_variance(self._wrap_model(model), *args, **kwargs) 92 | 93 | def training_losses( 94 | self, model, *args, **kwargs 95 | ): # pylint: disable=signature-differs 96 | return super().training_losses(self._wrap_model(model), *args, **kwargs) 97 | 98 | def condition_mean(self, cond_fn, *args, **kwargs): 99 | return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs) 100 | 101 | def condition_score(self, cond_fn, *args, **kwargs): 102 | return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs) 103 | 104 | def get_eps(self, model, *args, **kwargs): 105 | return super().get_eps(self._wrap_model(model), *args, **kwargs) 106 | 107 | def _wrap_model(self, model): 108 | if isinstance(model, _WrappedModel): 109 | return model 110 | return _WrappedModel( 111 | model, self.timestep_map, self.rescale_timesteps, self.original_num_steps 112 | ) 113 | 114 | def _scale_timesteps(self, t): 115 | # Scaling is done by the wrapped model. 116 | return t 117 | 118 | 119 | class _WrappedModel: 120 | def __init__(self, model, timestep_map, rescale_timesteps, original_num_steps): 121 | self.model = model 122 | self.timestep_map = timestep_map 123 | self.rescale_timesteps = rescale_timesteps 124 | self.original_num_steps = original_num_steps 125 | 126 | def __call__(self, x, ts, **kwargs): 127 | ts = ts.float() 128 | frac = ts.frac() 129 | map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype) 130 | new_ts_1 = map_tensor[ts.floor().long()] 131 | new_ts_2 = map_tensor[ts.ceil().long()] 132 | new_ts = th.lerp(new_ts_1, new_ts_2, frac) 133 | return self.model(x, new_ts, **kwargs) 134 | -------------------------------------------------------------------------------- /guided_diffusion/script_util.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import inspect 3 | 4 | from . import gaussian_diffusion as gd 5 | from .respace import SpacedDiffusion, space_timesteps 6 | from .unet import UNetModel 7 | 8 | NUM_CLASSES = 2 9 | 10 | 11 | def diffusion_defaults(): 12 | """ 13 | Defaults for image and classifier training. 14 | """ 15 | return dict( 16 | learn_sigma=False, 17 | diffusion_steps=1000, 18 | noise_schedule="linear", 19 | timestep_respacing="", 20 | use_kl=False, 21 | predict_xstart=False, 22 | rescale_timesteps=False, 23 | rescale_learned_sigmas=False, 24 | ) 25 | 26 | 27 | def classifier_defaults(): 28 | """ 29 | Defaults for classifier models. 30 | """ 31 | return dict( 32 | image_size=64, 33 | classifier_use_fp16=False, 34 | classifier_width=128, 35 | classifier_depth=2, 36 | classifier_attention_resolutions="32,16,8", # 16 37 | classifier_use_scale_shift_norm=True, # False 38 | classifier_resblock_updown=True, # False 39 | classifier_pool="attention", 40 | ) 41 | 42 | 43 | def model_and_diffusion_defaults(): 44 | """ 45 | Defaults for image training. 46 | """ 47 | res = dict( 48 | image_size=64, 49 | num_channels=128, 50 | num_res_blocks=2, 51 | num_heads=4, 52 | num_heads_upsample=-1, 53 | num_head_channels=-1, 54 | attention_resolutions="16,8", 55 | channel_mult="", 56 | dropout=0.0, 57 | class_cond=False, 58 | use_checkpoint=False, 59 | use_scale_shift_norm=True, 60 | resblock_updown=False, 61 | use_fp16=False, 62 | use_spatial_transformer=True, 63 | context_dim=1280, 64 | clip_embed_dim=None, 65 | image_condition=False, 66 | super_res_condition=False, 67 | ) 68 | res.update(diffusion_defaults()) 69 | return res 70 | 71 | 72 | def classifier_and_diffusion_defaults(): 73 | res = classifier_defaults() 74 | res.update(diffusion_defaults()) 75 | return res 76 | 77 | 78 | def create_model_and_diffusion( 79 | image_size, 80 | class_cond, 81 | learn_sigma, 82 | num_channels, 83 | num_res_blocks, 84 | channel_mult, 85 | num_heads, 86 | num_head_channels, 87 | num_heads_upsample, 88 | attention_resolutions, 89 | dropout, 90 | diffusion_steps, 91 | noise_schedule, 92 | timestep_respacing, 93 | use_kl, 94 | predict_xstart, 95 | rescale_timesteps, 96 | rescale_learned_sigmas, 97 | use_checkpoint, 98 | use_scale_shift_norm, 99 | resblock_updown, 100 | use_fp16, 101 | use_spatial_transformer, 102 | context_dim, 103 | clip_embed_dim, 104 | image_condition, 105 | super_res_condition, 106 | ): 107 | model = create_model( 108 | image_size, 109 | num_channels, 110 | num_res_blocks, 111 | channel_mult=channel_mult, 112 | learn_sigma=learn_sigma, 113 | class_cond=class_cond, 114 | use_checkpoint=use_checkpoint, 115 | attention_resolutions=attention_resolutions, 116 | num_heads=num_heads, 117 | num_head_channels=num_head_channels, 118 | num_heads_upsample=num_heads_upsample, 119 | use_scale_shift_norm=use_scale_shift_norm, 120 | dropout=dropout, 121 | resblock_updown=resblock_updown, 122 | use_fp16=use_fp16, 123 | use_spatial_transformer=use_spatial_transformer, 124 | context_dim=context_dim, 125 | clip_embed_dim=clip_embed_dim, 126 | image_condition=image_condition, 127 | super_res_condition=super_res_condition, 128 | ) 129 | diffusion = create_gaussian_diffusion( 130 | steps=diffusion_steps, 131 | learn_sigma=learn_sigma, 132 | noise_schedule=noise_schedule, 133 | use_kl=use_kl, 134 | predict_xstart=predict_xstart, 135 | rescale_timesteps=rescale_timesteps, 136 | rescale_learned_sigmas=rescale_learned_sigmas, 137 | timestep_respacing=timestep_respacing, 138 | ) 139 | return model, diffusion 140 | 141 | 142 | def create_model( 143 | image_size, 144 | num_channels, 145 | num_res_blocks, 146 | channel_mult="", 147 | learn_sigma=False, 148 | class_cond=False, 149 | use_checkpoint=False, 150 | attention_resolutions="16", 151 | num_heads=1, 152 | num_head_channels=-1, 153 | num_heads_upsample=-1, 154 | use_scale_shift_norm=False, 155 | dropout=0, 156 | resblock_updown=False, 157 | use_fp16=False, 158 | use_spatial_transformer=True, 159 | context_dim=1280, 160 | clip_embed_dim=None, 161 | image_condition=False, 162 | super_res_condition=False, 163 | ): 164 | if channel_mult == "": 165 | if image_size == 512: 166 | channel_mult = (0.5, 1, 1, 2, 2, 4, 4) 167 | elif image_size == 256: 168 | channel_mult = (1, 1, 2, 2, 4, 4) 169 | elif image_size == 128: 170 | channel_mult = (1, 1, 2, 3, 4) 171 | elif image_size == 64: 172 | channel_mult = (1, 2, 3, 4) 173 | elif image_size == 32: 174 | channel_mult = (1, 2, 4, 4) 175 | else: 176 | raise ValueError(f"unsupported image size: {image_size}") 177 | else: 178 | channel_mult = tuple(int(ch_mult) for ch_mult in channel_mult.split(",")) 179 | 180 | attention_ds = [] 181 | for res in attention_resolutions.split(","): 182 | attention_ds.append(image_size // int(res)) 183 | 184 | in_channels = 4 if image_size == 32 else 3 185 | out_channels = 4 186 | 187 | return UNetModel( 188 | image_size=image_size, 189 | in_channels=in_channels, 190 | model_channels=num_channels, 191 | out_channels=out_channels, 192 | num_res_blocks=num_res_blocks, 193 | attention_resolutions=tuple(attention_ds), 194 | dropout=dropout, 195 | channel_mult=channel_mult, 196 | num_classes=(NUM_CLASSES if class_cond else None), 197 | use_checkpoint=use_checkpoint, 198 | use_fp16=use_fp16, 199 | num_heads=num_heads, 200 | num_head_channels=num_head_channels, 201 | num_heads_upsample=num_heads_upsample, 202 | use_scale_shift_norm=use_scale_shift_norm, 203 | resblock_updown=resblock_updown, 204 | use_spatial_transformer=use_spatial_transformer, 205 | context_dim=context_dim, 206 | clip_embed_dim=clip_embed_dim, 207 | image_condition=image_condition, 208 | super_res_condition=super_res_condition, 209 | ) 210 | 211 | 212 | def create_classifier_and_diffusion( 213 | image_size, 214 | classifier_use_fp16, 215 | classifier_width, 216 | classifier_depth, 217 | classifier_attention_resolutions, 218 | classifier_use_scale_shift_norm, 219 | classifier_resblock_updown, 220 | classifier_pool, 221 | learn_sigma, 222 | diffusion_steps, 223 | noise_schedule, 224 | timestep_respacing, 225 | use_kl, 226 | predict_xstart, 227 | rescale_timesteps, 228 | rescale_learned_sigmas, 229 | ): 230 | classifier = create_classifier( 231 | image_size, 232 | classifier_use_fp16, 233 | classifier_width, 234 | classifier_depth, 235 | classifier_attention_resolutions, 236 | classifier_use_scale_shift_norm, 237 | classifier_resblock_updown, 238 | classifier_pool, 239 | ) 240 | diffusion = create_gaussian_diffusion( 241 | steps=diffusion_steps, 242 | learn_sigma=learn_sigma, 243 | noise_schedule=noise_schedule, 244 | use_kl=use_kl, 245 | predict_xstart=predict_xstart, 246 | rescale_timesteps=rescale_timesteps, 247 | rescale_learned_sigmas=rescale_learned_sigmas, 248 | timestep_respacing=timestep_respacing, 249 | ) 250 | return classifier, diffusion 251 | 252 | 253 | def create_classifier( 254 | image_size, 255 | classifier_use_fp16, 256 | classifier_width, 257 | classifier_depth, 258 | classifier_attention_resolutions, 259 | classifier_use_scale_shift_norm, 260 | classifier_resblock_updown, 261 | classifier_pool, 262 | ): 263 | if image_size == 512: 264 | channel_mult = (0.5, 1, 1, 2, 2, 4, 4) 265 | elif image_size == 256: 266 | channel_mult = (1, 1, 2, 2, 4, 4) 267 | elif image_size == 128: 268 | channel_mult = (1, 1, 2, 3, 4) 269 | elif image_size == 64: 270 | channel_mult = (1, 2, 3, 4) 271 | else: 272 | raise ValueError(f"unsupported image size: {image_size}") 273 | 274 | attention_ds = [] 275 | for res in classifier_attention_resolutions.split(","): 276 | attention_ds.append(image_size // int(res)) 277 | 278 | return EncoderUNetModel( 279 | image_size=image_size, 280 | in_channels=3, 281 | model_channels=classifier_width, 282 | out_channels=1000, 283 | num_res_blocks=classifier_depth, 284 | attention_resolutions=tuple(attention_ds), 285 | channel_mult=channel_mult, 286 | use_fp16=classifier_use_fp16, 287 | num_head_channels=64, 288 | use_scale_shift_norm=classifier_use_scale_shift_norm, 289 | resblock_updown=classifier_resblock_updown, 290 | pool=classifier_pool, 291 | ) 292 | 293 | 294 | def sr_model_and_diffusion_defaults(): 295 | res = model_and_diffusion_defaults() 296 | res["large_size"] = 256 297 | res["small_size"] = 64 298 | arg_names = inspect.getfullargspec(sr_create_model_and_diffusion)[0] 299 | for k in res.copy().keys(): 300 | if k not in arg_names: 301 | del res[k] 302 | return res 303 | 304 | 305 | def sr_create_model_and_diffusion( 306 | large_size, 307 | small_size, 308 | class_cond, 309 | learn_sigma, 310 | num_channels, 311 | num_res_blocks, 312 | num_heads, 313 | num_head_channels, 314 | num_heads_upsample, 315 | attention_resolutions, 316 | dropout, 317 | diffusion_steps, 318 | noise_schedule, 319 | timestep_respacing, 320 | use_kl, 321 | predict_xstart, 322 | rescale_timesteps, 323 | rescale_learned_sigmas, 324 | use_checkpoint, 325 | use_scale_shift_norm, 326 | resblock_updown, 327 | use_fp16, 328 | ): 329 | model = sr_create_model( 330 | large_size, 331 | small_size, 332 | num_channels, 333 | num_res_blocks, 334 | learn_sigma=learn_sigma, 335 | class_cond=class_cond, 336 | use_checkpoint=use_checkpoint, 337 | attention_resolutions=attention_resolutions, 338 | num_heads=num_heads, 339 | num_head_channels=num_head_channels, 340 | num_heads_upsample=num_heads_upsample, 341 | use_scale_shift_norm=use_scale_shift_norm, 342 | dropout=dropout, 343 | resblock_updown=resblock_updown, 344 | use_fp16=use_fp16, 345 | ) 346 | diffusion = create_gaussian_diffusion( 347 | steps=diffusion_steps, 348 | learn_sigma=learn_sigma, 349 | noise_schedule=noise_schedule, 350 | use_kl=use_kl, 351 | predict_xstart=predict_xstart, 352 | rescale_timesteps=rescale_timesteps, 353 | rescale_learned_sigmas=rescale_learned_sigmas, 354 | timestep_respacing=timestep_respacing, 355 | ) 356 | return model, diffusion 357 | 358 | 359 | def sr_create_model( 360 | large_size, 361 | small_size, 362 | num_channels, 363 | num_res_blocks, 364 | learn_sigma, 365 | class_cond, 366 | use_checkpoint, 367 | attention_resolutions, 368 | num_heads, 369 | num_head_channels, 370 | num_heads_upsample, 371 | use_scale_shift_norm, 372 | dropout, 373 | resblock_updown, 374 | use_fp16, 375 | ): 376 | _ = small_size # hack to prevent unused variable 377 | 378 | if large_size == 512: 379 | channel_mult = (1, 1, 2, 2, 4, 4) 380 | elif large_size == 256: 381 | channel_mult = (1, 1, 2, 2, 4, 4) 382 | elif large_size == 64: 383 | channel_mult = (1, 2, 3, 4) 384 | elif large_size == 32: 385 | channel_mult = (1, 2, 3, 4) 386 | else: 387 | raise ValueError(f"unsupported large size: {large_size}") 388 | 389 | attention_ds = [] 390 | for res in attention_resolutions.split(","): 391 | attention_ds.append(large_size // int(res)) 392 | 393 | return SuperResModel( 394 | image_size=large_size, 395 | in_channels=3, 396 | model_channels=num_channels, 397 | out_channels=(3 if not learn_sigma else 6), 398 | num_res_blocks=num_res_blocks, 399 | attention_resolutions=tuple(attention_ds), 400 | dropout=dropout, 401 | channel_mult=channel_mult, 402 | num_classes=(NUM_CLASSES if class_cond else None), 403 | use_checkpoint=use_checkpoint, 404 | num_heads=num_heads, 405 | num_head_channels=num_head_channels, 406 | num_heads_upsample=num_heads_upsample, 407 | use_scale_shift_norm=use_scale_shift_norm, 408 | resblock_updown=resblock_updown, 409 | use_fp16=use_fp16, 410 | ) 411 | 412 | 413 | def create_gaussian_diffusion( 414 | *, 415 | steps=1000, 416 | learn_sigma=False, 417 | sigma_small=False, 418 | noise_schedule="linear", 419 | use_kl=False, 420 | predict_xstart=False, 421 | rescale_timesteps=False, 422 | rescale_learned_sigmas=False, 423 | timestep_respacing="", 424 | ): 425 | betas = gd.get_named_beta_schedule(noise_schedule, steps) 426 | if use_kl: 427 | loss_type = gd.LossType.RESCALED_KL 428 | elif rescale_learned_sigmas: 429 | loss_type = gd.LossType.RESCALED_MSE 430 | else: 431 | loss_type = gd.LossType.MSE 432 | if not timestep_respacing: 433 | timestep_respacing = [steps] 434 | return SpacedDiffusion( 435 | use_timesteps=space_timesteps(steps, timestep_respacing), 436 | betas=betas, 437 | model_mean_type=( 438 | gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X 439 | ), 440 | model_var_type=( 441 | ( 442 | gd.ModelVarType.FIXED_LARGE 443 | if not sigma_small 444 | else gd.ModelVarType.FIXED_SMALL 445 | ) 446 | if not learn_sigma 447 | else gd.ModelVarType.LEARNED_RANGE 448 | ), 449 | loss_type=loss_type, 450 | rescale_timesteps=rescale_timesteps, 451 | ) 452 | 453 | 454 | def add_dict_to_argparser(parser, default_dict): 455 | for k, v in default_dict.items(): 456 | v_type = type(v) 457 | if v is None: 458 | v_type = str 459 | elif isinstance(v, bool): 460 | v_type = str2bool 461 | parser.add_argument(f"--{k}", default=v, type=v_type) 462 | 463 | 464 | def args_to_dict(args, keys): 465 | return {k: getattr(args, k) for k in keys} 466 | 467 | 468 | def str2bool(v): 469 | """ 470 | https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse 471 | """ 472 | if isinstance(v, bool): 473 | return v 474 | if v.lower() in ("yes", "true", "t", "y", "1"): 475 | return True 476 | elif v.lower() in ("no", "false", "f", "n", "0"): 477 | return False 478 | else: 479 | raise argparse.ArgumentTypeError("boolean value expected") 480 | -------------------------------------------------------------------------------- /ldm/models/diffusion/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/ldm/models/diffusion/__init__.py -------------------------------------------------------------------------------- /ldm/models/diffusion/classifier.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import pytorch_lightning as pl 4 | from omegaconf import OmegaConf 5 | from torch.nn import functional as F 6 | from torch.optim import AdamW 7 | from torch.optim.lr_scheduler import LambdaLR 8 | from copy import deepcopy 9 | from einops import rearrange 10 | from glob import glob 11 | from natsort import natsorted 12 | 13 | from ldm.modules.diffusionmodules.openaimodel import EncoderUNetModel, UNetModel 14 | from ldm.util import log_txt_as_img, default, ismap, instantiate_from_config 15 | 16 | __models__ = { 17 | 'class_label': EncoderUNetModel, 18 | 'segmentation': UNetModel 19 | } 20 | 21 | 22 | def disabled_train(self, mode=True): 23 | """Overwrite model.train with this function to make sure train/eval mode 24 | does not change anymore.""" 25 | return self 26 | 27 | 28 | class NoisyLatentImageClassifier(pl.LightningModule): 29 | 30 | def __init__(self, 31 | diffusion_path, 32 | num_classes, 33 | ckpt_path=None, 34 | pool='attention', 35 | label_key=None, 36 | diffusion_ckpt_path=None, 37 | scheduler_config=None, 38 | weight_decay=1.e-2, 39 | log_steps=10, 40 | monitor='val/loss', 41 | *args, 42 | **kwargs): 43 | super().__init__(*args, **kwargs) 44 | self.num_classes = num_classes 45 | # get latest config of diffusion model 46 | diffusion_config = natsorted(glob(os.path.join(diffusion_path, 'configs', '*-project.yaml')))[-1] 47 | self.diffusion_config = OmegaConf.load(diffusion_config).model 48 | self.diffusion_config.params.ckpt_path = diffusion_ckpt_path 49 | self.load_diffusion() 50 | 51 | self.monitor = monitor 52 | self.numd = self.diffusion_model.first_stage_model.encoder.num_resolutions - 1 53 | self.log_time_interval = self.diffusion_model.num_timesteps // log_steps 54 | self.log_steps = log_steps 55 | 56 | self.label_key = label_key if not hasattr(self.diffusion_model, 'cond_stage_key') \ 57 | else self.diffusion_model.cond_stage_key 58 | 59 | assert self.label_key is not None, 'label_key neither in diffusion model nor in model.params' 60 | 61 | if self.label_key not in __models__: 62 | raise NotImplementedError() 63 | 64 | self.load_classifier(ckpt_path, pool) 65 | 66 | self.scheduler_config = scheduler_config 67 | self.use_scheduler = self.scheduler_config is not None 68 | self.weight_decay = weight_decay 69 | 70 | def init_from_ckpt(self, path, ignore_keys=list(), only_model=False): 71 | sd = torch.load(path, map_location="cpu") 72 | if "state_dict" in list(sd.keys()): 73 | sd = sd["state_dict"] 74 | keys = list(sd.keys()) 75 | for k in keys: 76 | for ik in ignore_keys: 77 | if k.startswith(ik): 78 | print("Deleting key {} from state_dict.".format(k)) 79 | del sd[k] 80 | missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict( 81 | sd, strict=False) 82 | print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys") 83 | if len(missing) > 0: 84 | print(f"Missing Keys: {missing}") 85 | if len(unexpected) > 0: 86 | print(f"Unexpected Keys: {unexpected}") 87 | 88 | def load_diffusion(self): 89 | model = instantiate_from_config(self.diffusion_config) 90 | self.diffusion_model = model.eval() 91 | self.diffusion_model.train = disabled_train 92 | for param in self.diffusion_model.parameters(): 93 | param.requires_grad = False 94 | 95 | def load_classifier(self, ckpt_path, pool): 96 | model_config = deepcopy(self.diffusion_config.params.unet_config.params) 97 | model_config.in_channels = self.diffusion_config.params.unet_config.params.out_channels 98 | model_config.out_channels = self.num_classes 99 | if self.label_key == 'class_label': 100 | model_config.pool = pool 101 | 102 | self.model = __models__[self.label_key](**model_config) 103 | if ckpt_path is not None: 104 | print('#####################################################################') 105 | print(f'load from ckpt "{ckpt_path}"') 106 | print('#####################################################################') 107 | self.init_from_ckpt(ckpt_path) 108 | 109 | @torch.no_grad() 110 | def get_x_noisy(self, x, t, noise=None): 111 | noise = default(noise, lambda: torch.randn_like(x)) 112 | continuous_sqrt_alpha_cumprod = None 113 | if self.diffusion_model.use_continuous_noise: 114 | continuous_sqrt_alpha_cumprod = self.diffusion_model.sample_continuous_noise_level(x.shape[0], t + 1) 115 | # todo: make sure t+1 is correct here 116 | 117 | return self.diffusion_model.q_sample(x_start=x, t=t, noise=noise, 118 | continuous_sqrt_alpha_cumprod=continuous_sqrt_alpha_cumprod) 119 | 120 | def forward(self, x_noisy, t, *args, **kwargs): 121 | return self.model(x_noisy, t) 122 | 123 | @torch.no_grad() 124 | def get_input(self, batch, k): 125 | x = batch[k] 126 | if len(x.shape) == 3: 127 | x = x[..., None] 128 | x = rearrange(x, 'b h w c -> b c h w') 129 | x = x.to(memory_format=torch.contiguous_format).float() 130 | return x 131 | 132 | @torch.no_grad() 133 | def get_conditioning(self, batch, k=None): 134 | if k is None: 135 | k = self.label_key 136 | assert k is not None, 'Needs to provide label key' 137 | 138 | targets = batch[k].to(self.device) 139 | 140 | if self.label_key == 'segmentation': 141 | targets = rearrange(targets, 'b h w c -> b c h w') 142 | for down in range(self.numd): 143 | h, w = targets.shape[-2:] 144 | targets = F.interpolate(targets, size=(h // 2, w // 2), mode='nearest') 145 | 146 | # targets = rearrange(targets,'b c h w -> b h w c') 147 | 148 | return targets 149 | 150 | def compute_top_k(self, logits, labels, k, reduction="mean"): 151 | _, top_ks = torch.topk(logits, k, dim=1) 152 | if reduction == "mean": 153 | return (top_ks == labels[:, None]).float().sum(dim=-1).mean().item() 154 | elif reduction == "none": 155 | return (top_ks == labels[:, None]).float().sum(dim=-1) 156 | 157 | def on_train_epoch_start(self): 158 | # save some memory 159 | self.diffusion_model.model.to('cpu') 160 | 161 | @torch.no_grad() 162 | def write_logs(self, loss, logits, targets): 163 | log_prefix = 'train' if self.training else 'val' 164 | log = {} 165 | log[f"{log_prefix}/loss"] = loss.mean() 166 | log[f"{log_prefix}/acc@1"] = self.compute_top_k( 167 | logits, targets, k=1, reduction="mean" 168 | ) 169 | log[f"{log_prefix}/acc@5"] = self.compute_top_k( 170 | logits, targets, k=5, reduction="mean" 171 | ) 172 | 173 | self.log_dict(log, prog_bar=False, logger=True, on_step=self.training, on_epoch=True) 174 | self.log('loss', log[f"{log_prefix}/loss"], prog_bar=True, logger=False) 175 | self.log('global_step', self.global_step, logger=False, on_epoch=False, prog_bar=True) 176 | lr = self.optimizers().param_groups[0]['lr'] 177 | self.log('lr_abs', lr, on_step=True, logger=True, on_epoch=False, prog_bar=True) 178 | 179 | def shared_step(self, batch, t=None): 180 | x, *_ = self.diffusion_model.get_input(batch, k=self.diffusion_model.first_stage_key) 181 | targets = self.get_conditioning(batch) 182 | if targets.dim() == 4: 183 | targets = targets.argmax(dim=1) 184 | if t is None: 185 | t = torch.randint(0, self.diffusion_model.num_timesteps, (x.shape[0],), device=self.device).long() 186 | else: 187 | t = torch.full(size=(x.shape[0],), fill_value=t, device=self.device).long() 188 | x_noisy = self.get_x_noisy(x, t) 189 | logits = self(x_noisy, t) 190 | 191 | loss = F.cross_entropy(logits, targets, reduction='none') 192 | 193 | self.write_logs(loss.detach(), logits.detach(), targets.detach()) 194 | 195 | loss = loss.mean() 196 | return loss, logits, x_noisy, targets 197 | 198 | def training_step(self, batch, batch_idx): 199 | loss, *_ = self.shared_step(batch) 200 | return loss 201 | 202 | def reset_noise_accs(self): 203 | self.noisy_acc = {t: {'acc@1': [], 'acc@5': []} for t in 204 | range(0, self.diffusion_model.num_timesteps, self.diffusion_model.log_every_t)} 205 | 206 | def on_validation_start(self): 207 | self.reset_noise_accs() 208 | 209 | @torch.no_grad() 210 | def validation_step(self, batch, batch_idx): 211 | loss, *_ = self.shared_step(batch) 212 | 213 | for t in self.noisy_acc: 214 | _, logits, _, targets = self.shared_step(batch, t) 215 | self.noisy_acc[t]['acc@1'].append(self.compute_top_k(logits, targets, k=1, reduction='mean')) 216 | self.noisy_acc[t]['acc@5'].append(self.compute_top_k(logits, targets, k=5, reduction='mean')) 217 | 218 | return loss 219 | 220 | def configure_optimizers(self): 221 | optimizer = AdamW(self.model.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay) 222 | 223 | if self.use_scheduler: 224 | scheduler = instantiate_from_config(self.scheduler_config) 225 | 226 | print("Setting up LambdaLR scheduler...") 227 | scheduler = [ 228 | { 229 | 'scheduler': LambdaLR(optimizer, lr_lambda=scheduler.schedule), 230 | 'interval': 'step', 231 | 'frequency': 1 232 | }] 233 | return [optimizer], scheduler 234 | 235 | return optimizer 236 | 237 | @torch.no_grad() 238 | def log_images(self, batch, N=8, *args, **kwargs): 239 | log = dict() 240 | x = self.get_input(batch, self.diffusion_model.first_stage_key) 241 | log['inputs'] = x 242 | 243 | y = self.get_conditioning(batch) 244 | 245 | if self.label_key == 'class_label': 246 | y = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"]) 247 | log['labels'] = y 248 | 249 | if ismap(y): 250 | log['labels'] = self.diffusion_model.to_rgb(y) 251 | 252 | for step in range(self.log_steps): 253 | current_time = step * self.log_time_interval 254 | 255 | _, logits, x_noisy, _ = self.shared_step(batch, t=current_time) 256 | 257 | log[f'inputs@t{current_time}'] = x_noisy 258 | 259 | pred = F.one_hot(logits.argmax(dim=1), num_classes=self.num_classes) 260 | pred = rearrange(pred, 'b h w c -> b c h w') 261 | 262 | log[f'pred@t{current_time}'] = self.diffusion_model.to_rgb(pred) 263 | 264 | for key in log: 265 | log[key] = log[key][:N] 266 | 267 | return log 268 | -------------------------------------------------------------------------------- /ldm/models/diffusion/ddim.py: -------------------------------------------------------------------------------- 1 | """SAMPLING ONLY.""" 2 | 3 | import torch 4 | import numpy as np 5 | from tqdm import tqdm 6 | from functools import partial 7 | 8 | from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like 9 | 10 | 11 | class DDIMSampler(object): 12 | def __init__(self, model, schedule="linear", **kwargs): 13 | super().__init__() 14 | self.model = model 15 | self.ddpm_num_timesteps = model.num_timesteps 16 | self.schedule = schedule 17 | 18 | def register_buffer(self, name, attr): 19 | if type(attr) == torch.Tensor: 20 | if attr.device != torch.device("cuda"): 21 | attr = attr.to(torch.device("cuda")) 22 | setattr(self, name, attr) 23 | 24 | def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True): 25 | self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps, 26 | num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose) 27 | alphas_cumprod = self.model.alphas_cumprod 28 | assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep' 29 | to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device) 30 | 31 | self.register_buffer('betas', to_torch(self.model.betas)) 32 | self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod)) 33 | self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev)) 34 | 35 | # calculations for diffusion q(x_t | x_{t-1}) and others 36 | self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu()))) 37 | self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu()))) 38 | self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu()))) 39 | self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu()))) 40 | self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1))) 41 | 42 | # ddim sampling parameters 43 | ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(), 44 | ddim_timesteps=self.ddim_timesteps, 45 | eta=ddim_eta,verbose=verbose) 46 | self.register_buffer('ddim_sigmas', ddim_sigmas) 47 | self.register_buffer('ddim_alphas', ddim_alphas) 48 | self.register_buffer('ddim_alphas_prev', ddim_alphas_prev) 49 | self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas)) 50 | sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt( 51 | (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * ( 52 | 1 - self.alphas_cumprod / self.alphas_cumprod_prev)) 53 | self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps) 54 | 55 | @torch.no_grad() 56 | def sample(self, 57 | S, 58 | batch_size, 59 | shape, 60 | conditioning=None, 61 | callback=None, 62 | normals_sequence=None, 63 | img_callback=None, 64 | quantize_x0=False, 65 | eta=0., 66 | mask=None, 67 | x0=None, 68 | temperature=1., 69 | noise_dropout=0., 70 | score_corrector=None, 71 | corrector_kwargs=None, 72 | verbose=True, 73 | x_T=None, 74 | log_every_t=100, 75 | unconditional_guidance_scale=1., 76 | unconditional_conditioning=None, 77 | # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ... 78 | **kwargs 79 | ): 80 | if conditioning is not None: 81 | if isinstance(conditioning, dict): 82 | cbs = conditioning[list(conditioning.keys())[0]].shape[0] 83 | if cbs != batch_size: 84 | print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}") 85 | else: 86 | if conditioning.shape[0] != batch_size: 87 | print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}") 88 | 89 | self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose) 90 | # sampling 91 | C, H, W = shape 92 | size = (batch_size, C, H, W) 93 | print(f'Data shape for DDIM sampling is {size}, eta {eta}') 94 | 95 | samples, intermediates = self.ddim_sampling(conditioning, size, 96 | callback=callback, 97 | img_callback=img_callback, 98 | quantize_denoised=quantize_x0, 99 | mask=mask, x0=x0, 100 | ddim_use_original_steps=False, 101 | noise_dropout=noise_dropout, 102 | temperature=temperature, 103 | score_corrector=score_corrector, 104 | corrector_kwargs=corrector_kwargs, 105 | x_T=x_T, 106 | log_every_t=log_every_t, 107 | unconditional_guidance_scale=unconditional_guidance_scale, 108 | unconditional_conditioning=unconditional_conditioning, 109 | ) 110 | return samples, intermediates 111 | 112 | @torch.no_grad() 113 | def ddim_sampling(self, cond, shape, 114 | x_T=None, ddim_use_original_steps=False, 115 | callback=None, timesteps=None, quantize_denoised=False, 116 | mask=None, x0=None, img_callback=None, log_every_t=100, 117 | temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None, 118 | unconditional_guidance_scale=1., unconditional_conditioning=None,): 119 | device = self.model.betas.device 120 | b = shape[0] 121 | if x_T is None: 122 | img = torch.randn(shape, device=device) 123 | else: 124 | img = x_T 125 | 126 | if timesteps is None: 127 | timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps 128 | elif timesteps is not None and not ddim_use_original_steps: 129 | subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1 130 | timesteps = self.ddim_timesteps[:subset_end] 131 | 132 | intermediates = {'x_inter': [img], 'pred_x0': [img]} 133 | time_range = reversed(range(0,timesteps)) if ddim_use_original_steps else np.flip(timesteps) 134 | total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0] 135 | print(f"Running DDIM Sampling with {total_steps} timesteps") 136 | 137 | iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps) 138 | 139 | for i, step in enumerate(iterator): 140 | index = total_steps - i - 1 141 | ts = torch.full((b,), step, device=device, dtype=torch.long) 142 | 143 | if mask is not None: 144 | assert x0 is not None 145 | img_orig = self.model.q_sample(x0, ts) # TODO: deterministic forward pass? 146 | img = img_orig * mask + (1. - mask) * img 147 | 148 | outs = self.p_sample_ddim(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps, 149 | quantize_denoised=quantize_denoised, temperature=temperature, 150 | noise_dropout=noise_dropout, score_corrector=score_corrector, 151 | corrector_kwargs=corrector_kwargs, 152 | unconditional_guidance_scale=unconditional_guidance_scale, 153 | unconditional_conditioning=unconditional_conditioning) 154 | img, pred_x0 = outs 155 | if callback: callback(i) 156 | if img_callback: img_callback(pred_x0, i) 157 | 158 | if index % log_every_t == 0 or index == total_steps - 1: 159 | intermediates['x_inter'].append(img) 160 | intermediates['pred_x0'].append(pred_x0) 161 | 162 | return img, intermediates 163 | 164 | @torch.no_grad() 165 | def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False, 166 | temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None, 167 | unconditional_guidance_scale=1., unconditional_conditioning=None): 168 | b, *_, device = *x.shape, x.device 169 | 170 | if unconditional_conditioning is None or unconditional_guidance_scale == 1.: 171 | e_t = self.model.apply_model(x, t, c) 172 | else: 173 | x_in = torch.cat([x] * 2) 174 | t_in = torch.cat([t] * 2) 175 | c_in = torch.cat([unconditional_conditioning, c]) 176 | e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2) 177 | e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond) 178 | 179 | if score_corrector is not None: 180 | assert self.model.parameterization == "eps" 181 | e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs) 182 | 183 | alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas 184 | alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev 185 | sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas 186 | sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas 187 | # select parameters corresponding to the currently considered timestep 188 | a_t = torch.full((b, 1, 1, 1), alphas[index], device=device) 189 | a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device) 190 | sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device) 191 | sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device) 192 | 193 | # current prediction for x_0 194 | pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt() 195 | if quantize_denoised: 196 | pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0) 197 | # direction pointing to x_t 198 | dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t 199 | noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature 200 | if noise_dropout > 0.: 201 | noise = torch.nn.functional.dropout(noise, p=noise_dropout) 202 | x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise 203 | return x_prev, pred_x0 204 | -------------------------------------------------------------------------------- /ldm/models/diffusion/plms.py: -------------------------------------------------------------------------------- 1 | """SAMPLING ONLY.""" 2 | 3 | import torch 4 | import numpy as np 5 | from tqdm import tqdm 6 | from functools import partial 7 | 8 | from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like 9 | 10 | 11 | class PLMSSampler(object): 12 | def __init__(self, model, schedule="linear", **kwargs): 13 | super().__init__() 14 | self.model = model 15 | self.ddpm_num_timesteps = model.num_timesteps 16 | self.schedule = schedule 17 | 18 | def register_buffer(self, name, attr): 19 | if type(attr) == torch.Tensor: 20 | if attr.device != torch.device("cuda"): 21 | attr = attr.to(torch.device("cuda")) 22 | setattr(self, name, attr) 23 | 24 | def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True): 25 | if ddim_eta != 0: 26 | raise ValueError('ddim_eta must be 0 for PLMS') 27 | self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps, 28 | num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose) 29 | alphas_cumprod = self.model.alphas_cumprod 30 | assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep' 31 | to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device) 32 | 33 | self.register_buffer('betas', to_torch(self.model.betas)) 34 | self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod)) 35 | self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev)) 36 | 37 | # calculations for diffusion q(x_t | x_{t-1}) and others 38 | self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu()))) 39 | self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu()))) 40 | self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu()))) 41 | self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu()))) 42 | self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1))) 43 | 44 | # ddim sampling parameters 45 | ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(), 46 | ddim_timesteps=self.ddim_timesteps, 47 | eta=ddim_eta,verbose=verbose) 48 | self.register_buffer('ddim_sigmas', ddim_sigmas) 49 | self.register_buffer('ddim_alphas', ddim_alphas) 50 | self.register_buffer('ddim_alphas_prev', ddim_alphas_prev) 51 | self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas)) 52 | sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt( 53 | (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * ( 54 | 1 - self.alphas_cumprod / self.alphas_cumprod_prev)) 55 | self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps) 56 | 57 | @torch.no_grad() 58 | def sample(self, 59 | S, 60 | batch_size, 61 | shape, 62 | conditioning=None, 63 | callback=None, 64 | normals_sequence=None, 65 | img_callback=None, 66 | quantize_x0=False, 67 | eta=0., 68 | mask=None, 69 | x0=None, 70 | temperature=1., 71 | noise_dropout=0., 72 | score_corrector=None, 73 | corrector_kwargs=None, 74 | verbose=True, 75 | x_T=None, 76 | log_every_t=100, 77 | unconditional_guidance_scale=1., 78 | unconditional_conditioning=None, 79 | # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ... 80 | **kwargs 81 | ): 82 | if conditioning is not None: 83 | if isinstance(conditioning, dict): 84 | cbs = conditioning[list(conditioning.keys())[0]].shape[0] 85 | if cbs != batch_size: 86 | print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}") 87 | else: 88 | if conditioning.shape[0] != batch_size: 89 | print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}") 90 | 91 | self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose) 92 | # sampling 93 | C, H, W = shape 94 | size = (batch_size, C, H, W) 95 | print(f'Data shape for PLMS sampling is {size}') 96 | 97 | samples, intermediates = self.plms_sampling(conditioning, size, 98 | callback=callback, 99 | img_callback=img_callback, 100 | quantize_denoised=quantize_x0, 101 | mask=mask, x0=x0, 102 | ddim_use_original_steps=False, 103 | noise_dropout=noise_dropout, 104 | temperature=temperature, 105 | score_corrector=score_corrector, 106 | corrector_kwargs=corrector_kwargs, 107 | x_T=x_T, 108 | log_every_t=log_every_t, 109 | unconditional_guidance_scale=unconditional_guidance_scale, 110 | unconditional_conditioning=unconditional_conditioning, 111 | ) 112 | return samples, intermediates 113 | 114 | @torch.no_grad() 115 | def plms_sampling(self, cond, shape, 116 | x_T=None, ddim_use_original_steps=False, 117 | callback=None, timesteps=None, quantize_denoised=False, 118 | mask=None, x0=None, img_callback=None, log_every_t=100, 119 | temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None, 120 | unconditional_guidance_scale=1., unconditional_conditioning=None,): 121 | device = self.model.betas.device 122 | b = shape[0] 123 | if x_T is None: 124 | img = torch.randn(shape, device=device) 125 | else: 126 | img = x_T 127 | 128 | if timesteps is None: 129 | timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps 130 | elif timesteps is not None and not ddim_use_original_steps: 131 | subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1 132 | timesteps = self.ddim_timesteps[:subset_end] 133 | 134 | intermediates = {'x_inter': [img], 'pred_x0': [img]} 135 | time_range = list(reversed(range(0,timesteps))) if ddim_use_original_steps else np.flip(timesteps) 136 | total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0] 137 | print(f"Running PLMS Sampling with {total_steps} timesteps") 138 | 139 | iterator = tqdm(time_range, desc='PLMS Sampler', total=total_steps) 140 | old_eps = [] 141 | 142 | for i, step in enumerate(iterator): 143 | index = total_steps - i - 1 144 | ts = torch.full((b,), step, device=device, dtype=torch.long) 145 | ts_next = torch.full((b,), time_range[min(i + 1, len(time_range) - 1)], device=device, dtype=torch.long) 146 | 147 | if mask is not None: 148 | assert x0 is not None 149 | img_orig = self.model.q_sample(x0, ts) # TODO: deterministic forward pass? 150 | img = img_orig * mask + (1. - mask) * img 151 | 152 | outs = self.p_sample_plms(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps, 153 | quantize_denoised=quantize_denoised, temperature=temperature, 154 | noise_dropout=noise_dropout, score_corrector=score_corrector, 155 | corrector_kwargs=corrector_kwargs, 156 | unconditional_guidance_scale=unconditional_guidance_scale, 157 | unconditional_conditioning=unconditional_conditioning, 158 | old_eps=old_eps, t_next=ts_next) 159 | img, pred_x0, e_t = outs 160 | old_eps.append(e_t) 161 | if len(old_eps) >= 4: 162 | old_eps.pop(0) 163 | if callback: callback(i) 164 | if img_callback: img_callback(pred_x0, i) 165 | 166 | if index % log_every_t == 0 or index == total_steps - 1: 167 | intermediates['x_inter'].append(img) 168 | intermediates['pred_x0'].append(pred_x0) 169 | 170 | return img, intermediates 171 | 172 | @torch.no_grad() 173 | def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False, 174 | temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None, 175 | unconditional_guidance_scale=1., unconditional_conditioning=None, old_eps=None, t_next=None): 176 | b, *_, device = *x.shape, x.device 177 | 178 | def get_model_output(x, t): 179 | if unconditional_conditioning is None or unconditional_guidance_scale == 1.: 180 | e_t = self.model.apply_model(x, t, c) 181 | else: 182 | x_in = torch.cat([x] * 2) 183 | t_in = torch.cat([t] * 2) 184 | c_in = torch.cat([unconditional_conditioning, c]) 185 | e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2) 186 | e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond) 187 | 188 | if score_corrector is not None: 189 | assert self.model.parameterization == "eps" 190 | e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs) 191 | 192 | return e_t 193 | 194 | alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas 195 | alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev 196 | sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas 197 | sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas 198 | 199 | def get_x_prev_and_pred_x0(e_t, index): 200 | # select parameters corresponding to the currently considered timestep 201 | a_t = torch.full((b, 1, 1, 1), alphas[index], device=device) 202 | a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device) 203 | sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device) 204 | sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device) 205 | 206 | # current prediction for x_0 207 | pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt() 208 | if quantize_denoised: 209 | pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0) 210 | # direction pointing to x_t 211 | dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t 212 | noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature 213 | if noise_dropout > 0.: 214 | noise = torch.nn.functional.dropout(noise, p=noise_dropout) 215 | x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise 216 | return x_prev, pred_x0 217 | 218 | e_t = get_model_output(x, t) 219 | if len(old_eps) == 0: 220 | # Pseudo Improved Euler (2nd order) 221 | x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t, index) 222 | e_t_next = get_model_output(x_prev, t_next) 223 | e_t_prime = (e_t + e_t_next) / 2 224 | elif len(old_eps) == 1: 225 | # 2nd order Pseudo Linear Multistep (Adams-Bashforth) 226 | e_t_prime = (3 * e_t - old_eps[-1]) / 2 227 | elif len(old_eps) == 2: 228 | # 3nd order Pseudo Linear Multistep (Adams-Bashforth) 229 | e_t_prime = (23 * e_t - 16 * old_eps[-1] + 5 * old_eps[-2]) / 12 230 | elif len(old_eps) >= 3: 231 | # 4nd order Pseudo Linear Multistep (Adams-Bashforth) 232 | e_t_prime = (55 * e_t - 59 * old_eps[-1] + 37 * old_eps[-2] - 9 * old_eps[-3]) / 24 233 | 234 | x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t_prime, index) 235 | 236 | return x_prev, pred_x0, e_t 237 | -------------------------------------------------------------------------------- /ldm/modules/attention.py: -------------------------------------------------------------------------------- 1 | from inspect import isfunction 2 | import math 3 | import torch 4 | import torch.nn.functional as F 5 | from torch import nn, einsum 6 | from einops import rearrange, repeat 7 | 8 | from ldm.modules.diffusionmodules.util import checkpoint 9 | 10 | 11 | def exists(val): 12 | return val is not None 13 | 14 | 15 | def uniq(arr): 16 | return{el: True for el in arr}.keys() 17 | 18 | 19 | def default(val, d): 20 | if exists(val): 21 | return val 22 | return d() if isfunction(d) else d 23 | 24 | 25 | def max_neg_value(t): 26 | return -torch.finfo(t.dtype).max 27 | 28 | 29 | def init_(tensor): 30 | dim = tensor.shape[-1] 31 | std = 1 / math.sqrt(dim) 32 | tensor.uniform_(-std, std) 33 | return tensor 34 | 35 | 36 | # feedforward 37 | class GEGLU(nn.Module): 38 | def __init__(self, dim_in, dim_out): 39 | super().__init__() 40 | self.proj = nn.Linear(dim_in, dim_out * 2) 41 | 42 | def forward(self, x): 43 | x, gate = self.proj(x).chunk(2, dim=-1) 44 | return x * F.gelu(gate) 45 | 46 | 47 | class FeedForward(nn.Module): 48 | def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.): 49 | super().__init__() 50 | inner_dim = int(dim * mult) 51 | dim_out = default(dim_out, dim) 52 | project_in = nn.Sequential( 53 | nn.Linear(dim, inner_dim), 54 | nn.GELU() 55 | ) if not glu else GEGLU(dim, inner_dim) 56 | 57 | self.net = nn.Sequential( 58 | project_in, 59 | nn.Dropout(dropout), 60 | nn.Linear(inner_dim, dim_out) 61 | ) 62 | 63 | def forward(self, x): 64 | return self.net(x) 65 | 66 | 67 | def zero_module(module): 68 | """ 69 | Zero out the parameters of a module and return it. 70 | """ 71 | for p in module.parameters(): 72 | p.detach().zero_() 73 | return module 74 | 75 | 76 | def Normalize(in_channels): 77 | return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) 78 | 79 | 80 | class LinearAttention(nn.Module): 81 | def __init__(self, dim, heads=4, dim_head=32): 82 | super().__init__() 83 | self.heads = heads 84 | hidden_dim = dim_head * heads 85 | self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False) 86 | self.to_out = nn.Conv2d(hidden_dim, dim, 1) 87 | 88 | def forward(self, x): 89 | b, c, h, w = x.shape 90 | qkv = self.to_qkv(x) 91 | q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3) 92 | k = k.softmax(dim=-1) 93 | context = torch.einsum('bhdn,bhen->bhde', k, v) 94 | out = torch.einsum('bhde,bhdn->bhen', context, q) 95 | out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w) 96 | return self.to_out(out) 97 | 98 | 99 | class SpatialSelfAttention(nn.Module): 100 | def __init__(self, in_channels): 101 | super().__init__() 102 | self.in_channels = in_channels 103 | 104 | self.norm = Normalize(in_channels) 105 | self.q = torch.nn.Conv2d(in_channels, 106 | in_channels, 107 | kernel_size=1, 108 | stride=1, 109 | padding=0) 110 | self.k = torch.nn.Conv2d(in_channels, 111 | in_channels, 112 | kernel_size=1, 113 | stride=1, 114 | padding=0) 115 | self.v = torch.nn.Conv2d(in_channels, 116 | in_channels, 117 | kernel_size=1, 118 | stride=1, 119 | padding=0) 120 | self.proj_out = torch.nn.Conv2d(in_channels, 121 | in_channels, 122 | kernel_size=1, 123 | stride=1, 124 | padding=0) 125 | 126 | def forward(self, x): 127 | h_ = x 128 | h_ = self.norm(h_) 129 | q = self.q(h_) 130 | k = self.k(h_) 131 | v = self.v(h_) 132 | 133 | # compute attention 134 | b,c,h,w = q.shape 135 | q = rearrange(q, 'b c h w -> b (h w) c') 136 | k = rearrange(k, 'b c h w -> b c (h w)') 137 | w_ = torch.einsum('bij,bjk->bik', q, k) 138 | 139 | w_ = w_ * (int(c)**(-0.5)) 140 | w_ = torch.nn.functional.softmax(w_, dim=2) 141 | 142 | # attend to values 143 | v = rearrange(v, 'b c h w -> b c (h w)') 144 | w_ = rearrange(w_, 'b i j -> b j i') 145 | h_ = torch.einsum('bij,bjk->bik', v, w_) 146 | h_ = rearrange(h_, 'b c (h w) -> b c h w', h=h) 147 | h_ = self.proj_out(h_) 148 | 149 | return x+h_ 150 | 151 | 152 | class CrossAttention(nn.Module): 153 | def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.): 154 | super().__init__() 155 | inner_dim = dim_head * heads 156 | context_dim = default(context_dim, query_dim) 157 | 158 | self.scale = dim_head ** -0.5 159 | self.heads = heads 160 | 161 | self.to_q = nn.Linear(query_dim, inner_dim, bias=False) 162 | self.to_k = nn.Linear(context_dim, inner_dim, bias=False) 163 | self.to_v = nn.Linear(context_dim, inner_dim, bias=False) 164 | 165 | self.to_out = nn.Sequential( 166 | nn.Linear(inner_dim, query_dim), 167 | nn.Dropout(dropout) 168 | ) 169 | 170 | def forward(self, x, context=None, mask=None): 171 | h = self.heads 172 | 173 | q = self.to_q(x) 174 | context = default(context, x) 175 | k = self.to_k(context) 176 | v = self.to_v(context) 177 | 178 | q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v)) 179 | 180 | sim = einsum('b i d, b j d -> b i j', q, k) * self.scale 181 | 182 | if exists(mask): 183 | mask = rearrange(mask, 'b ... -> b (...)') 184 | max_neg_value = -torch.finfo(sim.dtype).max 185 | mask = repeat(mask, 'b j -> (b h) () j', h=h) 186 | sim.masked_fill_(~mask, max_neg_value) 187 | 188 | # attention, what we cannot get enough of 189 | attn = sim.softmax(dim=-1) 190 | 191 | out = einsum('b i j, b j d -> b i d', attn, v) 192 | out = rearrange(out, '(b h) n d -> b n (h d)', h=h) 193 | return self.to_out(out) 194 | 195 | 196 | class BasicTransformerBlock(nn.Module): 197 | def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True): 198 | super().__init__() 199 | self.attn1 = CrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout) # is a self-attention 200 | self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff) 201 | self.attn2 = CrossAttention(query_dim=dim, context_dim=context_dim, 202 | heads=n_heads, dim_head=d_head, dropout=dropout) # is self-attn if context is none 203 | self.norm1 = nn.LayerNorm(dim) 204 | self.norm2 = nn.LayerNorm(dim) 205 | self.norm3 = nn.LayerNorm(dim) 206 | self.checkpoint = checkpoint 207 | 208 | def forward(self, x, context=None): 209 | return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint) 210 | 211 | def _forward(self, x, context=None): 212 | x = self.attn1(self.norm1(x)) + x 213 | x = self.attn2(self.norm2(x), context=context) + x 214 | x = self.ff(self.norm3(x)) + x 215 | return x 216 | 217 | 218 | class SpatialTransformer(nn.Module): 219 | """ 220 | Transformer block for image-like data. 221 | First, project the input (aka embedding) 222 | and reshape to b, t, d. 223 | Then apply standard transformer action. 224 | Finally, reshape to image 225 | """ 226 | def __init__(self, in_channels, n_heads, d_head, 227 | depth=1, dropout=0., context_dim=None): 228 | super().__init__() 229 | self.in_channels = in_channels 230 | inner_dim = n_heads * d_head 231 | self.norm = Normalize(in_channels) 232 | 233 | self.proj_in = nn.Conv2d(in_channels, 234 | inner_dim, 235 | kernel_size=1, 236 | stride=1, 237 | padding=0) 238 | 239 | self.transformer_blocks = nn.ModuleList( 240 | [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim) 241 | for d in range(depth)] 242 | ) 243 | 244 | self.proj_out = zero_module(nn.Conv2d(inner_dim, 245 | in_channels, 246 | kernel_size=1, 247 | stride=1, 248 | padding=0)) 249 | 250 | def forward(self, x, context=None): 251 | # note: if no context is given, cross-attention defaults to self-attention 252 | b, c, h, w = x.shape 253 | x_in = x 254 | x = self.norm(x) 255 | x = self.proj_in(x) 256 | x = rearrange(x, 'b c h w -> b (h w) c') 257 | for block in self.transformer_blocks: 258 | x = block(x, context=context) 259 | x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w) 260 | x = self.proj_out(x) 261 | return x + x_in -------------------------------------------------------------------------------- /ldm/modules/diffusionmodules/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/ldm/modules/diffusionmodules/__init__.py -------------------------------------------------------------------------------- /ldm/modules/diffusionmodules/util.py: -------------------------------------------------------------------------------- 1 | # adopted from 2 | # https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py 3 | # and 4 | # https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py 5 | # and 6 | # https://github.com/openai/guided-diffusion/blob/0ba878e517b276c45d1195eb29f6f5f72659a05b/guided_diffusion/nn.py 7 | # 8 | # thanks! 9 | 10 | 11 | import os 12 | import math 13 | import torch 14 | import torch.nn as nn 15 | import numpy as np 16 | from einops import repeat 17 | 18 | from ldm.util import instantiate_from_config 19 | 20 | 21 | def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): 22 | if schedule == "linear": 23 | betas = ( 24 | torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 2 25 | ) 26 | 27 | elif schedule == "cosine": 28 | timesteps = ( 29 | torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s 30 | ) 31 | alphas = timesteps / (1 + cosine_s) * np.pi / 2 32 | alphas = torch.cos(alphas).pow(2) 33 | alphas = alphas / alphas[0] 34 | betas = 1 - alphas[1:] / alphas[:-1] 35 | betas = np.clip(betas, a_min=0, a_max=0.999) 36 | 37 | elif schedule == "sqrt_linear": 38 | betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) 39 | elif schedule == "sqrt": 40 | betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.5 41 | else: 42 | raise ValueError(f"schedule '{schedule}' unknown.") 43 | return betas.numpy() 44 | 45 | 46 | def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True): 47 | if ddim_discr_method == 'uniform': 48 | c = num_ddpm_timesteps // num_ddim_timesteps 49 | ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c))) 50 | elif ddim_discr_method == 'quad': 51 | ddim_timesteps = ((np.linspace(0, np.sqrt(num_ddpm_timesteps * .8), num_ddim_timesteps)) ** 2).astype(int) 52 | else: 53 | raise NotImplementedError(f'There is no ddim discretization method called "{ddim_discr_method}"') 54 | 55 | # assert ddim_timesteps.shape[0] == num_ddim_timesteps 56 | # add one to get the final alpha values right (the ones from first scale to data during sampling) 57 | steps_out = ddim_timesteps + 1 58 | if verbose: 59 | print(f'Selected timesteps for ddim sampler: {steps_out}') 60 | return steps_out 61 | 62 | 63 | def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta, verbose=True): 64 | # select alphas for computing the variance schedule 65 | alphas = alphacums[ddim_timesteps] 66 | alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist()) 67 | 68 | # according the the formula provided in https://arxiv.org/abs/2010.02502 69 | sigmas = eta * np.sqrt((1 - alphas_prev) / (1 - alphas) * (1 - alphas / alphas_prev)) 70 | if verbose: 71 | print(f'Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}') 72 | print(f'For the chosen value of eta, which is {eta}, ' 73 | f'this results in the following sigma_t schedule for ddim sampler {sigmas}') 74 | return sigmas, alphas, alphas_prev 75 | 76 | 77 | def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): 78 | """ 79 | Create a beta schedule that discretizes the given alpha_t_bar function, 80 | which defines the cumulative product of (1-beta) over time from t = [0,1]. 81 | :param num_diffusion_timesteps: the number of betas to produce. 82 | :param alpha_bar: a lambda that takes an argument t from 0 to 1 and 83 | produces the cumulative product of (1-beta) up to that 84 | part of the diffusion process. 85 | :param max_beta: the maximum beta to use; use values lower than 1 to 86 | prevent singularities. 87 | """ 88 | betas = [] 89 | for i in range(num_diffusion_timesteps): 90 | t1 = i / num_diffusion_timesteps 91 | t2 = (i + 1) / num_diffusion_timesteps 92 | betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta)) 93 | return np.array(betas) 94 | 95 | 96 | def extract_into_tensor(a, t, x_shape): 97 | b, *_ = t.shape 98 | out = a.gather(-1, t) 99 | return out.reshape(b, *((1,) * (len(x_shape) - 1))) 100 | 101 | 102 | def checkpoint(func, inputs, params, flag): 103 | """ 104 | Evaluate a function without caching intermediate activations, allowing for 105 | reduced memory at the expense of extra compute in the backward pass. 106 | :param func: the function to evaluate. 107 | :param inputs: the argument sequence to pass to `func`. 108 | :param params: a sequence of parameters `func` depends on but does not 109 | explicitly take as arguments. 110 | :param flag: if False, disable gradient checkpointing. 111 | """ 112 | if flag: 113 | args = tuple(inputs) + tuple(params) 114 | return CheckpointFunction.apply(func, len(inputs), *args) 115 | else: 116 | return func(*inputs) 117 | 118 | 119 | class CheckpointFunction(torch.autograd.Function): 120 | @staticmethod 121 | def forward(ctx, run_function, length, *args): 122 | ctx.run_function = run_function 123 | ctx.input_tensors = list(args[:length]) 124 | ctx.input_params = list(args[length:]) 125 | 126 | with torch.no_grad(): 127 | output_tensors = ctx.run_function(*ctx.input_tensors) 128 | return output_tensors 129 | 130 | @staticmethod 131 | def backward(ctx, *output_grads): 132 | ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors] 133 | with torch.enable_grad(): 134 | # Fixes a bug where the first op in run_function modifies the 135 | # Tensor storage in place, which is not allowed for detach()'d 136 | # Tensors. 137 | shallow_copies = [x.view_as(x) for x in ctx.input_tensors] 138 | output_tensors = ctx.run_function(*shallow_copies) 139 | input_grads = torch.autograd.grad( 140 | output_tensors, 141 | ctx.input_tensors + ctx.input_params, 142 | output_grads, 143 | allow_unused=True, 144 | ) 145 | del ctx.input_tensors 146 | del ctx.input_params 147 | del output_tensors 148 | return (None, None) + input_grads 149 | 150 | 151 | def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False): 152 | """ 153 | Create sinusoidal timestep embeddings. 154 | :param timesteps: a 1-D Tensor of N indices, one per batch element. 155 | These may be fractional. 156 | :param dim: the dimension of the output. 157 | :param max_period: controls the minimum frequency of the embeddings. 158 | :return: an [N x dim] Tensor of positional embeddings. 159 | """ 160 | if not repeat_only: 161 | half = dim // 2 162 | freqs = torch.exp( 163 | -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half 164 | ).to(device=timesteps.device) 165 | args = timesteps[:, None].float() * freqs[None] 166 | embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) 167 | if dim % 2: 168 | embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) 169 | else: 170 | embedding = repeat(timesteps, 'b -> b d', d=dim) 171 | return embedding 172 | 173 | 174 | def zero_module(module): 175 | """ 176 | Zero out the parameters of a module and return it. 177 | """ 178 | for p in module.parameters(): 179 | p.detach().zero_() 180 | return module 181 | 182 | 183 | def scale_module(module, scale): 184 | """ 185 | Scale the parameters of a module and return it. 186 | """ 187 | for p in module.parameters(): 188 | p.detach().mul_(scale) 189 | return module 190 | 191 | 192 | def mean_flat(tensor): 193 | """ 194 | Take the mean over all non-batch dimensions. 195 | """ 196 | return tensor.mean(dim=list(range(1, len(tensor.shape)))) 197 | 198 | 199 | def normalization(channels): 200 | """ 201 | Make a standard normalization layer. 202 | :param channels: number of input channels. 203 | :return: an nn.Module for normalization. 204 | """ 205 | return GroupNorm32(32, channels) 206 | 207 | 208 | # PyTorch 1.7 has SiLU, but we support PyTorch 1.5. 209 | class SiLU(nn.Module): 210 | def forward(self, x): 211 | return x * torch.sigmoid(x) 212 | 213 | 214 | class GroupNorm32(nn.GroupNorm): 215 | def forward(self, x): 216 | return super().forward(x.float()).type(x.dtype) 217 | 218 | def conv_nd(dims, *args, **kwargs): 219 | """ 220 | Create a 1D, 2D, or 3D convolution module. 221 | """ 222 | if dims == 1: 223 | return nn.Conv1d(*args, **kwargs) 224 | elif dims == 2: 225 | return nn.Conv2d(*args, **kwargs) 226 | elif dims == 3: 227 | return nn.Conv3d(*args, **kwargs) 228 | raise ValueError(f"unsupported dimensions: {dims}") 229 | 230 | 231 | def linear(*args, **kwargs): 232 | """ 233 | Create a linear module. 234 | """ 235 | return nn.Linear(*args, **kwargs) 236 | 237 | 238 | def avg_pool_nd(dims, *args, **kwargs): 239 | """ 240 | Create a 1D, 2D, or 3D average pooling module. 241 | """ 242 | if dims == 1: 243 | return nn.AvgPool1d(*args, **kwargs) 244 | elif dims == 2: 245 | return nn.AvgPool2d(*args, **kwargs) 246 | elif dims == 3: 247 | return nn.AvgPool3d(*args, **kwargs) 248 | raise ValueError(f"unsupported dimensions: {dims}") 249 | 250 | 251 | class HybridConditioner(nn.Module): 252 | 253 | def __init__(self, c_concat_config, c_crossattn_config): 254 | super().__init__() 255 | self.concat_conditioner = instantiate_from_config(c_concat_config) 256 | self.crossattn_conditioner = instantiate_from_config(c_crossattn_config) 257 | 258 | def forward(self, c_concat, c_crossattn): 259 | c_concat = self.concat_conditioner(c_concat) 260 | c_crossattn = self.crossattn_conditioner(c_crossattn) 261 | return {'c_concat': [c_concat], 'c_crossattn': [c_crossattn]} 262 | 263 | 264 | def noise_like(shape, device, repeat=False): 265 | repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1))) 266 | noise = lambda: torch.randn(shape, device=device) 267 | return repeat_noise() if repeat else noise() -------------------------------------------------------------------------------- /ldm/modules/distributions/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/ldm/modules/distributions/__init__.py -------------------------------------------------------------------------------- /ldm/modules/distributions/distributions.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | 4 | 5 | class AbstractDistribution: 6 | def sample(self): 7 | raise NotImplementedError() 8 | 9 | def mode(self): 10 | raise NotImplementedError() 11 | 12 | 13 | class DiracDistribution(AbstractDistribution): 14 | def __init__(self, value): 15 | self.value = value 16 | 17 | def sample(self): 18 | return self.value 19 | 20 | def mode(self): 21 | return self.value 22 | 23 | 24 | class DiagonalGaussianDistribution(object): 25 | def __init__(self, parameters, deterministic=False): 26 | self.parameters = parameters 27 | self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) 28 | self.logvar = torch.clamp(self.logvar, -30.0, 20.0) 29 | self.deterministic = deterministic 30 | self.std = torch.exp(0.5 * self.logvar) 31 | self.var = torch.exp(self.logvar) 32 | if self.deterministic: 33 | self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device) 34 | 35 | def sample(self): 36 | x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device) 37 | return x 38 | 39 | def kl(self, other=None): 40 | if self.deterministic: 41 | return torch.Tensor([0.]) 42 | else: 43 | if other is None: 44 | return 0.5 * torch.sum(torch.pow(self.mean, 2) 45 | + self.var - 1.0 - self.logvar, 46 | dim=[1, 2, 3]) 47 | else: 48 | return 0.5 * torch.sum( 49 | torch.pow(self.mean - other.mean, 2) / other.var 50 | + self.var / other.var - 1.0 - self.logvar + other.logvar, 51 | dim=[1, 2, 3]) 52 | 53 | def nll(self, sample, dims=[1,2,3]): 54 | if self.deterministic: 55 | return torch.Tensor([0.]) 56 | logtwopi = np.log(2.0 * np.pi) 57 | return 0.5 * torch.sum( 58 | logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, 59 | dim=dims) 60 | 61 | def mode(self): 62 | return self.mean 63 | 64 | 65 | def normal_kl(mean1, logvar1, mean2, logvar2): 66 | """ 67 | source: https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/losses.py#L12 68 | Compute the KL divergence between two gaussians. 69 | Shapes are automatically broadcasted, so batches can be compared to 70 | scalars, among other use cases. 71 | """ 72 | tensor = None 73 | for obj in (mean1, logvar1, mean2, logvar2): 74 | if isinstance(obj, torch.Tensor): 75 | tensor = obj 76 | break 77 | assert tensor is not None, "at least one argument must be a Tensor" 78 | 79 | # Force variances to be Tensors. Broadcasting helps convert scalars to 80 | # Tensors, but it does not work for torch.exp(). 81 | logvar1, logvar2 = [ 82 | x if isinstance(x, torch.Tensor) else torch.tensor(x).to(tensor) 83 | for x in (logvar1, logvar2) 84 | ] 85 | 86 | return 0.5 * ( 87 | -1.0 88 | + logvar2 89 | - logvar1 90 | + torch.exp(logvar1 - logvar2) 91 | + ((mean1 - mean2) ** 2) * torch.exp(-logvar2) 92 | ) 93 | -------------------------------------------------------------------------------- /ldm/modules/ema.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | 4 | 5 | class LitEma(nn.Module): 6 | def __init__(self, model, decay=0.9999, use_num_upates=True): 7 | super().__init__() 8 | if decay < 0.0 or decay > 1.0: 9 | raise ValueError('Decay must be between 0 and 1') 10 | 11 | self.m_name2s_name = {} 12 | self.register_buffer('decay', torch.tensor(decay, dtype=torch.float32)) 13 | self.register_buffer('num_updates', torch.tensor(0,dtype=torch.int) if use_num_upates 14 | else torch.tensor(-1,dtype=torch.int)) 15 | 16 | for name, p in model.named_parameters(): 17 | if p.requires_grad: 18 | #remove as '.'-character is not allowed in buffers 19 | s_name = name.replace('.','') 20 | self.m_name2s_name.update({name:s_name}) 21 | self.register_buffer(s_name,p.clone().detach().data) 22 | 23 | self.collected_params = [] 24 | 25 | def forward(self,model): 26 | decay = self.decay 27 | 28 | if self.num_updates >= 0: 29 | self.num_updates += 1 30 | decay = min(self.decay,(1 + self.num_updates) / (10 + self.num_updates)) 31 | 32 | one_minus_decay = 1.0 - decay 33 | 34 | with torch.no_grad(): 35 | m_param = dict(model.named_parameters()) 36 | shadow_params = dict(self.named_buffers()) 37 | 38 | for key in m_param: 39 | if m_param[key].requires_grad: 40 | sname = self.m_name2s_name[key] 41 | shadow_params[sname] = shadow_params[sname].type_as(m_param[key]) 42 | shadow_params[sname].sub_(one_minus_decay * (shadow_params[sname] - m_param[key])) 43 | else: 44 | assert not key in self.m_name2s_name 45 | 46 | def copy_to(self, model): 47 | m_param = dict(model.named_parameters()) 48 | shadow_params = dict(self.named_buffers()) 49 | for key in m_param: 50 | if m_param[key].requires_grad: 51 | m_param[key].data.copy_(shadow_params[self.m_name2s_name[key]].data) 52 | else: 53 | assert not key in self.m_name2s_name 54 | 55 | def store(self, parameters): 56 | """ 57 | Save the current parameters for restoring later. 58 | Args: 59 | parameters: Iterable of `torch.nn.Parameter`; the parameters to be 60 | temporarily stored. 61 | """ 62 | self.collected_params = [param.clone() for param in parameters] 63 | 64 | def restore(self, parameters): 65 | """ 66 | Restore the parameters stored with the `store` method. 67 | Useful to validate the model with EMA parameters without affecting the 68 | original optimization process. Store the parameters before the 69 | `copy_to` method. After validation (or model saving), use this to 70 | restore the former parameters. 71 | Args: 72 | parameters: Iterable of `torch.nn.Parameter`; the parameters to be 73 | updated with the stored parameters. 74 | """ 75 | for c_param, param in zip(self.collected_params, parameters): 76 | param.data.copy_(c_param.data) 77 | -------------------------------------------------------------------------------- /ldm/modules/encoders/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/ldm/modules/encoders/__init__.py -------------------------------------------------------------------------------- /ldm/modules/encoders/modules.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from functools import partial 4 | 5 | from ldm.modules.x_transformer import Encoder, TransformerWrapper # TODO: can we directly rely on lucidrains code and simply add this as a reuirement? --> test 6 | 7 | 8 | class AbstractEncoder(nn.Module): 9 | def __init__(self): 10 | super().__init__() 11 | 12 | def encode(self, *args, **kwargs): 13 | raise NotImplementedError 14 | 15 | 16 | 17 | class ClassEmbedder(nn.Module): 18 | def __init__(self, embed_dim, n_classes=1000, key='class'): 19 | super().__init__() 20 | self.key = key 21 | self.embedding = nn.Embedding(n_classes, embed_dim) 22 | 23 | def forward(self, batch, key=None): 24 | if key is None: 25 | key = self.key 26 | # this is for use in crossattn 27 | c = batch[key][:, None] 28 | c = self.embedding(c) 29 | return c 30 | 31 | 32 | class TransformerEmbedder(AbstractEncoder): 33 | """Some transformer encoder layers""" 34 | def __init__(self, n_embed, n_layer, vocab_size, max_seq_len=77, device="cuda"): 35 | super().__init__() 36 | self.device = device 37 | self.transformer = TransformerWrapper(num_tokens=vocab_size, max_seq_len=max_seq_len, 38 | attn_layers=Encoder(dim=n_embed, depth=n_layer)) 39 | 40 | def forward(self, tokens): 41 | tokens = tokens.to(self.device) # meh 42 | z = self.transformer(tokens, return_embeddings=True) 43 | return z 44 | 45 | def encode(self, x): 46 | return self(x) 47 | 48 | 49 | class BERTTokenizer(AbstractEncoder): 50 | """ Uses a pretrained BERT tokenizer by huggingface. Vocab size: 30522 (?)""" 51 | def __init__(self, device="cuda", vq_interface=True, max_length=77): 52 | super().__init__() 53 | from transformers import BertTokenizerFast # TODO: add to reuquirements 54 | self.tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased") 55 | self.device = device 56 | self.vq_interface = vq_interface 57 | self.max_length = max_length 58 | 59 | def forward(self, text): 60 | batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True, 61 | return_overflowing_tokens=False, padding="max_length", return_tensors="pt") 62 | tokens = batch_encoding["input_ids"].to(self.device) 63 | return tokens 64 | 65 | @torch.no_grad() 66 | def encode(self, text): 67 | tokens = self(text) 68 | if not self.vq_interface: 69 | return tokens 70 | return None, None, [None, None, tokens] 71 | 72 | def decode(self, text): 73 | return text 74 | 75 | 76 | class BERTEmbedder(AbstractEncoder): 77 | """Uses the BERT tokenizr model and add some transformer encoder layers""" 78 | def __init__(self, n_embed, n_layer, vocab_size=30522, max_seq_len=77, 79 | device="cuda",use_tokenizer=True, embedding_dropout=0.0): 80 | super().__init__() 81 | self.use_tknz_fn = use_tokenizer 82 | if self.use_tknz_fn: 83 | self.tknz_fn = BERTTokenizer(vq_interface=False, max_length=max_seq_len) 84 | self.device = device 85 | self.transformer = TransformerWrapper(num_tokens=vocab_size, max_seq_len=max_seq_len, 86 | attn_layers=Encoder(dim=n_embed, depth=n_layer), 87 | emb_dropout=embedding_dropout) 88 | 89 | def forward(self, text): 90 | if self.use_tknz_fn: 91 | tokens = self.tknz_fn(text)#.to(self.device) 92 | else: 93 | tokens = text 94 | z = self.transformer(tokens, return_embeddings=True) 95 | return z 96 | 97 | def encode(self, text): 98 | # output of length 77 99 | return self(text) 100 | 101 | 102 | class SpatialRescaler(nn.Module): 103 | def __init__(self, 104 | n_stages=1, 105 | method='bilinear', 106 | multiplier=0.5, 107 | in_channels=3, 108 | out_channels=None, 109 | bias=False): 110 | super().__init__() 111 | self.n_stages = n_stages 112 | assert self.n_stages >= 0 113 | assert method in ['nearest','linear','bilinear','trilinear','bicubic','area'] 114 | self.multiplier = multiplier 115 | self.interpolator = partial(torch.nn.functional.interpolate, mode=method) 116 | self.remap_output = out_channels is not None 117 | if self.remap_output: 118 | print(f'Spatial Rescaler mapping from {in_channels} to {out_channels} channels after resizing.') 119 | self.channel_mapper = nn.Conv2d(in_channels,out_channels,1,bias=bias) 120 | 121 | def forward(self,x): 122 | for stage in range(self.n_stages): 123 | x = self.interpolator(x, scale_factor=self.multiplier) 124 | 125 | 126 | if self.remap_output: 127 | x = self.channel_mapper(x) 128 | return x 129 | 130 | def encode(self, x): 131 | return self(x) 132 | -------------------------------------------------------------------------------- /ldm/modules/image_degradation/__init__.py: -------------------------------------------------------------------------------- 1 | from ldm.modules.image_degradation.bsrgan import degradation_bsrgan_variant as degradation_fn_bsr 2 | from ldm.modules.image_degradation.bsrgan_light import degradation_bsrgan_variant as degradation_fn_bsr_light 3 | -------------------------------------------------------------------------------- /ldm/modules/image_degradation/utils/test.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LAION-AI/ldm-finetune/0d25ab3e859827b2e0e109e2676ccecd99d63ad7/ldm/modules/image_degradation/utils/test.png -------------------------------------------------------------------------------- /ldm/modules/losses/__init__.py: -------------------------------------------------------------------------------- 1 | from ldm.modules.losses.contperceptual import LPIPSWithDiscriminator -------------------------------------------------------------------------------- /ldm/modules/losses/contperceptual.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from taming.modules.losses.vqperceptual import * # TODO: taming dependency yes/no? 5 | 6 | 7 | class LPIPSWithDiscriminator(nn.Module): 8 | def __init__(self, disc_start, logvar_init=0.0, kl_weight=1.0, pixelloss_weight=1.0, 9 | disc_num_layers=3, disc_in_channels=3, disc_factor=1.0, disc_weight=1.0, 10 | perceptual_weight=1.0, use_actnorm=False, disc_conditional=False, 11 | disc_loss="hinge"): 12 | 13 | super().__init__() 14 | assert disc_loss in ["hinge", "vanilla"] 15 | self.kl_weight = kl_weight 16 | self.pixel_weight = pixelloss_weight 17 | self.perceptual_loss = LPIPS().eval() 18 | self.perceptual_weight = perceptual_weight 19 | # output log variance 20 | self.logvar = nn.Parameter(torch.ones(size=()) * logvar_init) 21 | 22 | self.discriminator = NLayerDiscriminator(input_nc=disc_in_channels, 23 | n_layers=disc_num_layers, 24 | use_actnorm=use_actnorm 25 | ).apply(weights_init) 26 | self.discriminator_iter_start = disc_start 27 | self.disc_loss = hinge_d_loss if disc_loss == "hinge" else vanilla_d_loss 28 | self.disc_factor = disc_factor 29 | self.discriminator_weight = disc_weight 30 | self.disc_conditional = disc_conditional 31 | 32 | def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer=None): 33 | if last_layer is not None: 34 | nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0] 35 | g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0] 36 | else: 37 | nll_grads = torch.autograd.grad(nll_loss, self.last_layer[0], retain_graph=True)[0] 38 | g_grads = torch.autograd.grad(g_loss, self.last_layer[0], retain_graph=True)[0] 39 | 40 | d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4) 41 | d_weight = torch.clamp(d_weight, 0.0, 1e4).detach() 42 | d_weight = d_weight * self.discriminator_weight 43 | return d_weight 44 | 45 | def forward(self, inputs, reconstructions, posteriors, optimizer_idx, 46 | global_step, last_layer=None, cond=None, split="train", 47 | weights=None): 48 | rec_loss = torch.abs(inputs.contiguous() - reconstructions.contiguous()) 49 | if self.perceptual_weight > 0: 50 | p_loss = self.perceptual_loss(inputs.contiguous(), reconstructions.contiguous()) 51 | rec_loss = rec_loss + self.perceptual_weight * p_loss 52 | 53 | nll_loss = rec_loss / torch.exp(self.logvar) + self.logvar 54 | weighted_nll_loss = nll_loss 55 | if weights is not None: 56 | weighted_nll_loss = weights*nll_loss 57 | weighted_nll_loss = torch.sum(weighted_nll_loss) / weighted_nll_loss.shape[0] 58 | nll_loss = torch.sum(nll_loss) / nll_loss.shape[0] 59 | kl_loss = posteriors.kl() 60 | kl_loss = torch.sum(kl_loss) / kl_loss.shape[0] 61 | 62 | # now the GAN part 63 | if optimizer_idx == 0: 64 | # generator update 65 | if cond is None: 66 | assert not self.disc_conditional 67 | logits_fake = self.discriminator(reconstructions.contiguous()) 68 | else: 69 | assert self.disc_conditional 70 | logits_fake = self.discriminator(torch.cat((reconstructions.contiguous(), cond), dim=1)) 71 | g_loss = -torch.mean(logits_fake) 72 | 73 | if self.disc_factor > 0.0: 74 | try: 75 | d_weight = self.calculate_adaptive_weight(nll_loss, g_loss, last_layer=last_layer) 76 | except RuntimeError: 77 | assert not self.training 78 | d_weight = torch.tensor(0.0) 79 | else: 80 | d_weight = torch.tensor(0.0) 81 | 82 | disc_factor = adopt_weight(self.disc_factor, global_step, threshold=self.discriminator_iter_start) 83 | loss = weighted_nll_loss + self.kl_weight * kl_loss + d_weight * disc_factor * g_loss 84 | 85 | log = {"{}/total_loss".format(split): loss.clone().detach().mean(), "{}/logvar".format(split): self.logvar.detach(), 86 | "{}/kl_loss".format(split): kl_loss.detach().mean(), "{}/nll_loss".format(split): nll_loss.detach().mean(), 87 | "{}/rec_loss".format(split): rec_loss.detach().mean(), 88 | "{}/d_weight".format(split): d_weight.detach(), 89 | "{}/disc_factor".format(split): torch.tensor(disc_factor), 90 | "{}/g_loss".format(split): g_loss.detach().mean(), 91 | } 92 | return loss, log 93 | 94 | if optimizer_idx == 1: 95 | # second pass for discriminator update 96 | if cond is None: 97 | logits_real = self.discriminator(inputs.contiguous().detach()) 98 | logits_fake = self.discriminator(reconstructions.contiguous().detach()) 99 | else: 100 | logits_real = self.discriminator(torch.cat((inputs.contiguous().detach(), cond), dim=1)) 101 | logits_fake = self.discriminator(torch.cat((reconstructions.contiguous().detach(), cond), dim=1)) 102 | 103 | disc_factor = adopt_weight(self.disc_factor, global_step, threshold=self.discriminator_iter_start) 104 | d_loss = disc_factor * self.disc_loss(logits_real, logits_fake) 105 | 106 | log = {"{}/disc_loss".format(split): d_loss.clone().detach().mean(), 107 | "{}/logits_real".format(split): logits_real.detach().mean(), 108 | "{}/logits_fake".format(split): logits_fake.detach().mean() 109 | } 110 | return d_loss, log 111 | 112 | -------------------------------------------------------------------------------- /ldm/modules/losses/vqperceptual.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | import torch.nn.functional as F 4 | from einops import repeat 5 | 6 | from taming.modules.discriminator.model import NLayerDiscriminator, weights_init 7 | from taming.modules.losses.lpips import LPIPS 8 | from taming.modules.losses.vqperceptual import hinge_d_loss, vanilla_d_loss 9 | 10 | 11 | def hinge_d_loss_with_exemplar_weights(logits_real, logits_fake, weights): 12 | assert weights.shape[0] == logits_real.shape[0] == logits_fake.shape[0] 13 | loss_real = torch.mean(F.relu(1. - logits_real), dim=[1,2,3]) 14 | loss_fake = torch.mean(F.relu(1. + logits_fake), dim=[1,2,3]) 15 | loss_real = (weights * loss_real).sum() / weights.sum() 16 | loss_fake = (weights * loss_fake).sum() / weights.sum() 17 | d_loss = 0.5 * (loss_real + loss_fake) 18 | return d_loss 19 | 20 | def adopt_weight(weight, global_step, threshold=0, value=0.): 21 | if global_step < threshold: 22 | weight = value 23 | return weight 24 | 25 | 26 | def measure_perplexity(predicted_indices, n_embed): 27 | # src: https://github.com/karpathy/deep-vector-quantization/blob/main/model.py 28 | # eval cluster perplexity. when perplexity == num_embeddings then all clusters are used exactly equally 29 | encodings = F.one_hot(predicted_indices, n_embed).float().reshape(-1, n_embed) 30 | avg_probs = encodings.mean(0) 31 | perplexity = (-(avg_probs * torch.log(avg_probs + 1e-10)).sum()).exp() 32 | cluster_use = torch.sum(avg_probs > 0) 33 | return perplexity, cluster_use 34 | 35 | def l1(x, y): 36 | return torch.abs(x-y) 37 | 38 | 39 | def l2(x, y): 40 | return torch.pow((x-y), 2) 41 | 42 | 43 | class VQLPIPSWithDiscriminator(nn.Module): 44 | def __init__(self, disc_start, codebook_weight=1.0, pixelloss_weight=1.0, 45 | disc_num_layers=3, disc_in_channels=3, disc_factor=1.0, disc_weight=1.0, 46 | perceptual_weight=1.0, use_actnorm=False, disc_conditional=False, 47 | disc_ndf=64, disc_loss="hinge", n_classes=None, perceptual_loss="lpips", 48 | pixel_loss="l1"): 49 | super().__init__() 50 | assert disc_loss in ["hinge", "vanilla"] 51 | assert perceptual_loss in ["lpips", "clips", "dists"] 52 | assert pixel_loss in ["l1", "l2"] 53 | self.codebook_weight = codebook_weight 54 | self.pixel_weight = pixelloss_weight 55 | if perceptual_loss == "lpips": 56 | print(f"{self.__class__.__name__}: Running with LPIPS.") 57 | self.perceptual_loss = LPIPS().eval() 58 | else: 59 | raise ValueError(f"Unknown perceptual loss: >> {perceptual_loss} <<") 60 | self.perceptual_weight = perceptual_weight 61 | 62 | if pixel_loss == "l1": 63 | self.pixel_loss = l1 64 | else: 65 | self.pixel_loss = l2 66 | 67 | self.discriminator = NLayerDiscriminator(input_nc=disc_in_channels, 68 | n_layers=disc_num_layers, 69 | use_actnorm=use_actnorm, 70 | ndf=disc_ndf 71 | ).apply(weights_init) 72 | self.discriminator_iter_start = disc_start 73 | if disc_loss == "hinge": 74 | self.disc_loss = hinge_d_loss 75 | elif disc_loss == "vanilla": 76 | self.disc_loss = vanilla_d_loss 77 | else: 78 | raise ValueError(f"Unknown GAN loss '{disc_loss}'.") 79 | print(f"VQLPIPSWithDiscriminator running with {disc_loss} loss.") 80 | self.disc_factor = disc_factor 81 | self.discriminator_weight = disc_weight 82 | self.disc_conditional = disc_conditional 83 | self.n_classes = n_classes 84 | 85 | def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer=None): 86 | if last_layer is not None: 87 | nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0] 88 | g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0] 89 | else: 90 | nll_grads = torch.autograd.grad(nll_loss, self.last_layer[0], retain_graph=True)[0] 91 | g_grads = torch.autograd.grad(g_loss, self.last_layer[0], retain_graph=True)[0] 92 | 93 | d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4) 94 | d_weight = torch.clamp(d_weight, 0.0, 1e4).detach() 95 | d_weight = d_weight * self.discriminator_weight 96 | return d_weight 97 | 98 | def forward(self, codebook_loss, inputs, reconstructions, optimizer_idx, 99 | global_step, last_layer=None, cond=None, split="train", predicted_indices=None): 100 | if not exists(codebook_loss): 101 | codebook_loss = torch.tensor([0.]).to(inputs.device) 102 | #rec_loss = torch.abs(inputs.contiguous() - reconstructions.contiguous()) 103 | rec_loss = self.pixel_loss(inputs.contiguous(), reconstructions.contiguous()) 104 | if self.perceptual_weight > 0: 105 | p_loss = self.perceptual_loss(inputs.contiguous(), reconstructions.contiguous()) 106 | rec_loss = rec_loss + self.perceptual_weight * p_loss 107 | else: 108 | p_loss = torch.tensor([0.0]) 109 | 110 | nll_loss = rec_loss 111 | #nll_loss = torch.sum(nll_loss) / nll_loss.shape[0] 112 | nll_loss = torch.mean(nll_loss) 113 | 114 | # now the GAN part 115 | if optimizer_idx == 0: 116 | # generator update 117 | if cond is None: 118 | assert not self.disc_conditional 119 | logits_fake = self.discriminator(reconstructions.contiguous()) 120 | else: 121 | assert self.disc_conditional 122 | logits_fake = self.discriminator(torch.cat((reconstructions.contiguous(), cond), dim=1)) 123 | g_loss = -torch.mean(logits_fake) 124 | 125 | try: 126 | d_weight = self.calculate_adaptive_weight(nll_loss, g_loss, last_layer=last_layer) 127 | except RuntimeError: 128 | assert not self.training 129 | d_weight = torch.tensor(0.0) 130 | 131 | disc_factor = adopt_weight(self.disc_factor, global_step, threshold=self.discriminator_iter_start) 132 | loss = nll_loss + d_weight * disc_factor * g_loss + self.codebook_weight * codebook_loss.mean() 133 | 134 | log = {"{}/total_loss".format(split): loss.clone().detach().mean(), 135 | "{}/quant_loss".format(split): codebook_loss.detach().mean(), 136 | "{}/nll_loss".format(split): nll_loss.detach().mean(), 137 | "{}/rec_loss".format(split): rec_loss.detach().mean(), 138 | "{}/p_loss".format(split): p_loss.detach().mean(), 139 | "{}/d_weight".format(split): d_weight.detach(), 140 | "{}/disc_factor".format(split): torch.tensor(disc_factor), 141 | "{}/g_loss".format(split): g_loss.detach().mean(), 142 | } 143 | if predicted_indices is not None: 144 | assert self.n_classes is not None 145 | with torch.no_grad(): 146 | perplexity, cluster_usage = measure_perplexity(predicted_indices, self.n_classes) 147 | log[f"{split}/perplexity"] = perplexity 148 | log[f"{split}/cluster_usage"] = cluster_usage 149 | return loss, log 150 | 151 | if optimizer_idx == 1: 152 | # second pass for discriminator update 153 | if cond is None: 154 | logits_real = self.discriminator(inputs.contiguous().detach()) 155 | logits_fake = self.discriminator(reconstructions.contiguous().detach()) 156 | else: 157 | logits_real = self.discriminator(torch.cat((inputs.contiguous().detach(), cond), dim=1)) 158 | logits_fake = self.discriminator(torch.cat((reconstructions.contiguous().detach(), cond), dim=1)) 159 | 160 | disc_factor = adopt_weight(self.disc_factor, global_step, threshold=self.discriminator_iter_start) 161 | d_loss = disc_factor * self.disc_loss(logits_real, logits_fake) 162 | 163 | log = {"{}/disc_loss".format(split): d_loss.clone().detach().mean(), 164 | "{}/logits_real".format(split): logits_real.detach().mean(), 165 | "{}/logits_fake".format(split): logits_fake.detach().mean() 166 | } 167 | return d_loss, log 168 | -------------------------------------------------------------------------------- /ldm/util.py: -------------------------------------------------------------------------------- 1 | import importlib 2 | 3 | import torch 4 | import numpy as np 5 | 6 | from inspect import isfunction 7 | from PIL import Image, ImageDraw, ImageFont 8 | 9 | 10 | def log_txt_as_img(wh, xc, size=10): 11 | # wh a tuple of (width, height) 12 | # xc a list of captions to plot 13 | b = len(xc) 14 | txts = list() 15 | for bi in range(b): 16 | txt = Image.new("RGB", wh, color="white") 17 | draw = ImageDraw.Draw(txt) 18 | font = ImageFont.truetype('data/DejaVuSans.ttf', size=size) 19 | nc = int(40 * (wh[0] / 256)) 20 | lines = "\n".join(xc[bi][start:start + nc] for start in range(0, len(xc[bi]), nc)) 21 | 22 | try: 23 | draw.text((0, 0), lines, fill="black", font=font) 24 | except UnicodeEncodeError: 25 | print("Cant encode string for logging. Skipping.") 26 | 27 | txt = np.array(txt).transpose(2, 0, 1) / 127.5 - 1.0 28 | txts.append(txt) 29 | txts = np.stack(txts) 30 | txts = torch.tensor(txts) 31 | return txts 32 | 33 | 34 | def ismap(x): 35 | if not isinstance(x, torch.Tensor): 36 | return False 37 | return (len(x.shape) == 4) and (x.shape[1] > 3) 38 | 39 | 40 | def isimage(x): 41 | if not isinstance(x,torch.Tensor): 42 | return False 43 | return (len(x.shape) == 4) and (x.shape[1] == 3 or x.shape[1] == 1) 44 | 45 | 46 | def exists(x): 47 | return x is not None 48 | 49 | 50 | def default(val, d): 51 | if exists(val): 52 | return val 53 | return d() if isfunction(d) else d 54 | 55 | 56 | def mean_flat(tensor): 57 | """ 58 | https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/nn.py#L86 59 | Take the mean over all non-batch dimensions. 60 | """ 61 | return tensor.mean(dim=list(range(1, len(tensor.shape)))) 62 | 63 | 64 | def count_params(model, verbose=False): 65 | total_params = sum(p.numel() for p in model.parameters()) 66 | if verbose: 67 | print(f"{model.__class__.__name__} has {total_params*1.e-6:.2f} M params.") 68 | return total_params 69 | 70 | 71 | def instantiate_from_config(config): 72 | if not "target" in config: 73 | if config == '__is_first_stage__': 74 | return None 75 | elif config == "__is_unconditional__": 76 | return None 77 | raise KeyError("Expected key `target` to instantiate.") 78 | return get_obj_from_str(config["target"])(**config.get("params", dict())) 79 | 80 | 81 | def get_obj_from_str(string, reload=False): 82 | module, cls = string.rsplit(".", 1) 83 | if reload: 84 | module_imp = importlib.import_module(module) 85 | importlib.reload(module_imp) 86 | return getattr(importlib.import_module(module, package=None), cls) -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | axial-positional-embedding==0.2.1 2 | albumentations>=0.4.3 3 | blobfile>=1.2.9 4 | braceexpand>=0.1.7 5 | Cython==0.29.30 6 | DALL-E==0.1 7 | dalle-pytorch>=1.5.2 8 | ftfy>=5.8 9 | einops>=0.4.1 10 | huggingface-hub>=0.5.1 11 | imageio>=2.9.0 12 | imageio-ffmpeg>=0.4.2 13 | omegaconf>=2.1.2 14 | mpi4py>=3.1.3 15 | pytorch-lightning>=1.6 16 | PyYAML>=6.0 17 | regex>=2022.4.24 18 | rotary-embedding-torch>=0.1.5 19 | setuptools<=59.5.0 20 | tokenizers>=0.12.1 21 | torchmetrics>=0.8.0 22 | tqdm>=4.64.0 23 | transformers>=4.18.0 24 | torch-fidelity>=0.3.0 25 | wandb==0.12.17 -------------------------------------------------------------------------------- /sample_inpaint.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | from pathlib import Path 4 | 5 | from guided_diffusion.inpaint_util import (prepare_inpaint_models, 6 | sample_inpaint) 7 | 8 | os.environ[ 9 | "TOKENIZERS_PARALLELISM" 10 | ] = "false" # required to avoid errors with transformers lib 11 | 12 | 13 | def parse_args(): 14 | parser = argparse.ArgumentParser() 15 | parser.add_argument("--prompts", type=str, default="") 16 | parser.add_argument("--negative", type=str, default="") 17 | parser.add_argument("--init_image", type=str, default=None) 18 | parser.add_argument("--mask", type=str, default=None) 19 | parser.add_argument("--guidance_scale", type=float, default=5.0) 20 | parser.add_argument("--steps", type=int, default=100) 21 | parser.add_argument("--batch_size", type=int, default=1) 22 | parser.add_argument("--width", type=int, default=256) 23 | parser.add_argument("--height", type=int, default=256) 24 | parser.add_argument("--init_skip_fraction", type=float, default=0.0) 25 | parser.add_argument("--aesthetic_rating", type=int, default=9) 26 | parser.add_argument("--aesthetic_weight", type=float, default=0.5) 27 | parser.add_argument("--seed", type=int, default=0) 28 | parser.add_argument("--intermediate_outputs", type=bool, default=False) 29 | parser.add_argument("--model_path", type=str, default="inpaint.pt") 30 | parser.add_argument("--output_dir", type=str, default="inpaint_outputs") 31 | return parser.parse_args() 32 | 33 | 34 | if __name__ == "__main__": 35 | args = parse_args() 36 | prompts = args.prompts 37 | negative = args.negative 38 | init_image = args.init_image 39 | mask = args.mask 40 | guidance_scale = args.guidance_scale 41 | steps = args.steps 42 | batch_size = args.batch_size 43 | width = args.width 44 | height = args.height 45 | init_skip_fraction = args.init_skip_fraction 46 | aesthetic_rating = args.aesthetic_rating 47 | aesthetic_weight = args.aesthetic_weight 48 | seed = args.seed 49 | intermediate_outputs = args.intermediate_outputs 50 | model_path = args.model_path 51 | output_dir = args.output_dir 52 | 53 | inpaint_models = prepare_inpaint_models( 54 | inpaint_model_path=model_path, device="cuda", use_fp16=False 55 | ) 56 | 57 | if ".txt" in prompts and Path(prompts).exists(): 58 | with open(prompts, "r") as f: 59 | prompts = f.readlines() 60 | print(f"Read {len(prompts)} prompts from {prompts}") 61 | else: 62 | prompts = [prompts] 63 | 64 | for prompt in prompts: 65 | print(f"Generating prompt: {prompt}") 66 | generations = list( 67 | sample_inpaint( 68 | prompt=prompt, 69 | negative=negative, 70 | init_image=init_image, 71 | mask=mask, 72 | guidance_scale=guidance_scale, 73 | steps=steps, 74 | batch_size=batch_size, 75 | width=width, 76 | height=height, 77 | init_skip_fraction=init_skip_fraction, 78 | aesthetic_rating=aesthetic_rating, 79 | aesthetic_weight=aesthetic_weight, 80 | seed=seed, 81 | intermediate_outputs=intermediate_outputs, 82 | output_dir=output_dir, 83 | loaded_models=inpaint_models, 84 | ) 85 | ) 86 | -------------------------------------------------------------------------------- /scripts/image_train_inpaint.py: -------------------------------------------------------------------------------- 1 | """ 2 | Train a diffusion model on images. 3 | """ 4 | 5 | import argparse 6 | import random 7 | 8 | import torch 9 | from torchvision import transforms 10 | 11 | from encoders.modules import BERTEmbedder 12 | from guided_diffusion import dist_util, logger 13 | from guided_diffusion.image_text_datasets import load_data 14 | from guided_diffusion.resample import create_named_schedule_sampler 15 | from guided_diffusion.script_util import ( 16 | add_dict_to_argparser, 17 | args_to_dict, 18 | create_model_and_diffusion, 19 | model_and_diffusion_defaults, 20 | ) 21 | from guided_diffusion.train_util import TrainLoop 22 | 23 | def set_requires_grad(model, value): 24 | """ 25 | Set the requires_grad flag of all parameters in the model. 26 | """ 27 | for param in model.parameters(): 28 | param.requires_grad = value 29 | 30 | def main(): 31 | args = create_argparser().parse_args() 32 | 33 | 34 | dist_util.setup_dist() 35 | logger.configure() 36 | 37 | from dist.clip_custom import clip # make clip end up on the right device 38 | 39 | logger.log("loading clip...") 40 | clip_model, _ = clip.load('ViT-L/14', device=dist_util.dev(), jit=False) 41 | clip_model.eval().requires_grad_(False) 42 | set_requires_grad(clip_model, False) 43 | 44 | del clip_model.visual 45 | 46 | logger.log("loading vae...") 47 | 48 | encoder = torch.load(args.kl_model, map_location="cpu") 49 | if args.use_fp16: 50 | encoder = encoder.half() 51 | encoder.to(dist_util.dev()) 52 | encoder.eval() 53 | set_requires_grad(encoder, False) 54 | 55 | del encoder.loss 56 | 57 | logger.log("loading text encoder...") 58 | 59 | 60 | bert = BERTEmbedder(1280, 32) 61 | bert_state_dict = torch.load(args.bert_model, map_location="cpu") 62 | bert.load_state_dict(bert_state_dict) 63 | 64 | if args.use_fp16: 65 | bert = bert.half() 66 | bert = bert.to(dist_util.dev()) 67 | bert.eval() 68 | set_requires_grad(bert, False) 69 | 70 | diffusion_config = model_and_diffusion_defaults() 71 | logger.log("creating model and diffusion...") 72 | model, diffusion = create_model_and_diffusion( 73 | **args_to_dict(args, diffusion_config.keys()) 74 | ) 75 | 76 | model.to(dist_util.dev()) 77 | 78 | logger.log('total base parameters', sum(x.numel() for x in model.parameters())) 79 | 80 | schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion) 81 | 82 | logger.log("creating data loader...") 83 | 84 | data = load_latent_data( 85 | encoder=encoder, 86 | bert=bert, 87 | clip_model=clip_model, 88 | clip=clip, 89 | data_dir=args.data_dir, 90 | batch_size=args.batch_size, 91 | epochs=args.epochs, 92 | shard_size=args.shard_size, 93 | image_key=args.image_key, 94 | caption_key=args.caption_key, 95 | cache_dir=args.cache_dir, 96 | random_crop=args.random_crop, 97 | random_flip=args.random_flip, 98 | ) 99 | logger.log("training...") 100 | TrainLoop( 101 | model=model, 102 | bert=bert, 103 | diffusion=diffusion, 104 | diffusion_config=diffusion_config, 105 | kl_model=encoder, 106 | clip_model=clip_model, 107 | data=data, 108 | batch_size=args.batch_size, 109 | microbatch=args.microbatch, 110 | lr=args.lr, 111 | ema_rate=args.ema_rate, 112 | log_interval=args.log_interval, 113 | sample_interval=args.sample_interval, 114 | save_interval=args.save_interval, 115 | resume_checkpoint=args.resume_checkpoint, 116 | use_fp16=args.use_fp16, 117 | fp16_scale_growth=args.fp16_scale_growth, 118 | schedule_sampler=schedule_sampler, 119 | weight_decay=args.weight_decay, 120 | lr_anneal_steps=args.lr_anneal_steps, 121 | ).run_loop() 122 | 123 | 124 | def load_latent_data( 125 | encoder=None, 126 | bert=None, 127 | clip_model=None, 128 | clip=None, 129 | data_dir=None, 130 | batch_size=None, 131 | epochs=20, 132 | shard_size=10000, 133 | image_key="jpg", 134 | caption_key="txt", 135 | cache_dir="cache", 136 | random_crop=False, 137 | random_flip=False, 138 | use_fp16=False, 139 | ): 140 | data = load_data( 141 | data_dir=data_dir, 142 | batch_size=batch_size, 143 | random_crop=random_crop, 144 | random_flip=random_flip, 145 | image_key=image_key, 146 | caption_key=caption_key, 147 | cache_dir=cache_dir, 148 | epochs=epochs, 149 | shard_size=shard_size, # TODO 150 | ) 151 | blur = transforms.GaussianBlur(kernel_size=(15, 15), sigma=(0.1, 5)) 152 | for batch, text in data: 153 | batch = batch.to(dist_util.dev()) 154 | model_kwargs = {} 155 | 156 | text = list(text) 157 | for i in range(len(text)): 158 | if random.randint(0, 100) < 20: 159 | text[i] = "" 160 | 161 | text_emb = bert.encode(text).to(dist_util.dev()) 162 | 163 | clip_text = clip.tokenize(text, truncate=True).to(dist_util.dev()) 164 | clip_emb = clip_model.encode_text(clip_text) 165 | 166 | model_kwargs["context"] = text_emb.float() 167 | model_kwargs["clip_embed"] = clip_emb.float() 168 | 169 | batch = batch.to(dist_util.dev()) 170 | encoder_input = batch.half() if use_fp16 else batch.float() 171 | emb = encoder.encode(encoder_input).sample() 172 | if use_fp16: 173 | emb = emb.half() 174 | else: 175 | emb = emb.float() 176 | emb *= 0.18215 177 | 178 | emb_cond = emb.detach().clone() 179 | 180 | for i in range(batch.shape[0]): 181 | if random.randint(0, 100) < 20: 182 | emb_cond[i, :, :, :] = 0 # unconditional 183 | else: 184 | if random.randint(0, 100) < 50: 185 | mask = torch.randn(1, 32, 32).to(dist_util.dev()) 186 | mask = blur(mask) 187 | mask = mask > 0 188 | mask = mask.repeat(4, 1, 1) 189 | mask = mask.float() 190 | emb_cond[i] *= mask 191 | else: 192 | # mask out 4 random rectangles 193 | for j in range(random.randint(1, 4)): 194 | max_area = 32 * 16 195 | w = random.randint(1, 32) 196 | h = random.randint(1, 32) 197 | if w * h > max_area: 198 | if random.randint(0, 100) < 50: 199 | w = max_area // h 200 | else: 201 | h = max_area // w 202 | if w == 32: 203 | offsetx = 0 204 | else: 205 | offsetx = random.randint(0, 32 - w) 206 | if h == 32: 207 | offsety = 0 208 | else: 209 | offsety = random.randint(0, 32 - h) 210 | emb_cond[i, :, offsety : offsety + h, offsetx : offsetx + w] = 0 211 | 212 | model_kwargs["image_embed"] = emb_cond.float() 213 | 214 | yield emb, model_kwargs 215 | 216 | 217 | def create_argparser(): 218 | defaults = dict( 219 | data_dir="", 220 | schedule_sampler="uniform", 221 | lr=1e-4, 222 | weight_decay=0.0, 223 | lr_anneal_steps=0, 224 | batch_size=1, 225 | microbatch=-1, # -1 disables microbatches 226 | ema_rate="0.999", # comma-separated list of EMA values 227 | log_interval=10, 228 | save_interval=10000, 229 | sample_interval=100, 230 | resume_checkpoint="", 231 | use_fp16=False, 232 | fp16_scale_growth=1e-3, 233 | kl_model=None, 234 | bert_model=None, 235 | epochs=20, 236 | shard_size=10000, 237 | image_key="jpg", 238 | caption_key="txt", 239 | cache_dir="cache", 240 | random_crop=False, 241 | random_flip=False, 242 | ) 243 | defaults.update(model_and_diffusion_defaults()) 244 | 245 | defaults["clip_embed_dim"] = 768 246 | defaults["image_condition"] = True 247 | 248 | parser = argparse.ArgumentParser() 249 | add_dict_to_argparser(parser, defaults) 250 | return parser 251 | 252 | 253 | if __name__ == "__main__": 254 | main() 255 | -------------------------------------------------------------------------------- /scripts/image_train_latent.py: -------------------------------------------------------------------------------- 1 | """ 2 | Train a diffusion model on images. 3 | """ 4 | 5 | import argparse 6 | import random 7 | 8 | import torch 9 | 10 | from encoders.modules import BERTEmbedder 11 | from guided_diffusion import dist_util, logger 12 | from guided_diffusion.image_text_datasets import load_data 13 | from guided_diffusion.resample import create_named_schedule_sampler 14 | from guided_diffusion.script_util import (add_dict_to_argparser, args_to_dict, 15 | create_model_and_diffusion, 16 | model_and_diffusion_defaults) 17 | from guided_diffusion.train_util import TrainLoop 18 | 19 | 20 | def set_requires_grad(model, value): 21 | for param in model.parameters(): 22 | param.requires_grad = value 23 | 24 | 25 | def main(): 26 | args = create_argparser().parse_args() 27 | 28 | dist_util.setup_dist() 29 | logger.configure() 30 | 31 | logger.log("loading vae...") 32 | 33 | encoder = torch.load(args.kl_model, map_location="cpu") 34 | encoder.to(dist_util.dev()) 35 | encoder.eval() 36 | set_requires_grad(encoder, False) 37 | 38 | del encoder.decoder 39 | del encoder.loss 40 | 41 | logger.log("loading text encoder...") 42 | 43 | bert = BERTEmbedder(1280, 32) 44 | sd = torch.load(args.bert_model, map_location="cpu") 45 | bert.load_state_dict(sd) 46 | 47 | bert.to(dist_util.dev()) 48 | bert.eval() 49 | set_requires_grad(bert, False) 50 | 51 | logger.log("creating model and diffusion...") 52 | model, diffusion = create_model_and_diffusion( 53 | **args_to_dict(args, model_and_diffusion_defaults().keys()) 54 | ) 55 | 56 | model.to(dist_util.dev()) 57 | 58 | logger.log("total base parameters", sum(x.numel() for x in model.parameters())) 59 | 60 | schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion) 61 | 62 | logger.log("creating data loader...") 63 | data = load_latent_data( 64 | encoder, 65 | bert, 66 | data_dir=args.data_dir, 67 | batch_size=args.batch_size, 68 | image_size=args.image_size, 69 | ) 70 | logger.log("training...") 71 | TrainLoop( 72 | model=model, 73 | diffusion=diffusion, 74 | data=data, 75 | batch_size=args.batch_size, 76 | microbatch=args.microbatch, 77 | lr=args.lr, 78 | ema_rate=args.ema_rate, 79 | log_interval=args.log_interval, 80 | save_interval=args.save_interval, 81 | resume_checkpoint=args.resume_checkpoint, 82 | use_fp16=args.use_fp16, 83 | fp16_scale_growth=args.fp16_scale_growth, 84 | schedule_sampler=schedule_sampler, 85 | weight_decay=args.weight_decay, 86 | lr_anneal_steps=args.lr_anneal_steps, 87 | ).run_loop() 88 | 89 | 90 | def load_latent_data(encoder, bert, data_dir, batch_size, image_size): 91 | data = load_data( 92 | data_dir=data_dir, 93 | batch_size=batch_size, 94 | image_size=256, 95 | class_cond=False, 96 | ) 97 | for batch, model_kwargs, text in data: 98 | 99 | text = list(text) 100 | for i in range(len(text)): 101 | if random.randint(0, 100) < 20: 102 | text[i] = "" 103 | 104 | text_emb = bert.encode(text).to(dist_util.dev()).half() 105 | 106 | model_kwargs["context"] = text_emb 107 | 108 | batch = batch.to(dist_util.dev()) 109 | emb = encoder.encode(batch).sample().half() 110 | emb *= 0.18215 111 | 112 | yield emb, model_kwargs 113 | 114 | 115 | def create_argparser(): 116 | defaults = dict( 117 | data_dir="", 118 | schedule_sampler="uniform", 119 | lr=1e-4, 120 | weight_decay=0.0, 121 | lr_anneal_steps=0, 122 | batch_size=1, 123 | microbatch=-1, # -1 disables microbatches 124 | ema_rate="0.9999", # comma-separated list of EMA values 125 | log_interval=10, 126 | save_interval=10000, 127 | resume_checkpoint="", 128 | use_fp16=False, 129 | fp16_scale_growth=1e-3, 130 | kl_model=None, 131 | bert_model=None, 132 | ) 133 | defaults.update(model_and_diffusion_defaults()) 134 | parser = argparse.ArgumentParser() 135 | add_dict_to_argparser(parser, defaults) 136 | return parser 137 | 138 | 139 | if __name__ == "__main__": 140 | main() 141 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | setup( 4 | name="guided-diffusion", 5 | py_modules=["guided_diffusion"], 6 | install_requires=[ 7 | "blobfile>=1.0.5", 8 | "torch", 9 | "torchvision", 10 | "tqdm", 11 | "Cython" 12 | ], 13 | ) 14 | --------------------------------------------------------------------------------