├── LICENSE ├── README.md ├── datasets ├── README.md └── lsun_bedroom.py ├── encoders ├── __pycache__ │ ├── modules.cpython-38.pyc │ └── x_transformer.cpython-38.pyc ├── modules.py └── x_transformer.py ├── evaluations ├── README.md ├── evaluator.py └── requirements.txt ├── guided_diffusion.egg-info ├── PKG-INFO ├── SOURCES.txt ├── dependency_links.txt ├── requires.txt └── top_level.txt ├── guided_diffusion ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-38.pyc │ ├── dist_util.cpython-38.pyc │ ├── fp16_util.cpython-38.pyc │ ├── gaussian_diffusion.cpython-38.pyc │ ├── image_datasets_classifier.cpython-38.pyc │ ├── image_text_datasets.cpython-38.pyc │ ├── logger.cpython-38.pyc │ ├── losses.cpython-38.pyc │ ├── nn.cpython-38.pyc │ ├── resample.cpython-38.pyc │ ├── respace.cpython-38.pyc │ ├── script_util.cpython-38.pyc │ ├── train_util.cpython-38.pyc │ └── unet.cpython-38.pyc ├── dist_util.py ├── fp16_util.py ├── gaussian_diffusion.py ├── image_datasets.py ├── image_datasets_classifier.py ├── image_text_datasets.py ├── logger.py ├── losses.py ├── nn.py ├── resample.py ├── respace.py ├── script_util.py ├── train_util.py └── unet.py ├── kl.yaml ├── merge.py ├── model-card.md ├── sample.py ├── scripts ├── classifier_sample.py ├── classifier_train.py ├── classifier_train_stable.py ├── image_nll.py ├── image_sample.py ├── image_train.py ├── image_train_inpaint.py ├── image_train_stable.py ├── image_train_super.py ├── super_res_sample.py └── super_res_train.py ├── setup.py ├── split.py ├── super_large.py ├── train.sh ├── train_classifier.sh └── train_inpaint.sh /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 | # GLID-3-XL-stable 2 | 3 | GLID-3-xl-stable is [stable diffusion](https://github.com/CompVis/stable-diffusion) back-ported to the OpenAI guided diffusion codebase, for easier development and training. 4 | 5 | features: 6 | 7 | [inpainting/outpainting](https://github.com/Jack000/glid-3-xl-stable/wiki/Custom-inpainting-model) 8 | 9 | [classifier guided stable diffusion](https://github.com/Jack000/glid-3-xl-stable/wiki/Classifier-guided-stable-diffusion) 10 | 11 | [super-resolution](https://github.com/Jack000/glid-3-xl-stable/wiki/Double-diffusion-for-more-detailed-upscaling) 12 | 13 | # Install 14 | 15 | First install [latent diffusion](https://github.com/CompVis/latent-diffusion) 16 | ``` 17 | # then 18 | git clone https://github.com/Jack000/glid-3-xl-stable 19 | cd glid-3-xl-stable 20 | pip install -e . 21 | 22 | # install mpi and mpi4py for training 23 | sudo apt install libopenmpi-dev 24 | pip install mpi4py 25 | 26 | ``` 27 | 28 | # Get model files from stable diffusion 29 | 30 | ``` 31 | # split checkpoint 32 | python split.py sd-v1-4.ckpt 33 | 34 | # you should now have diffusion.pt and kl.pt 35 | 36 | # alternatively 37 | wget -O diffusion.pt https://huggingface.co/Jack000/glid-3-xl-stable/resolve/main/default/diffusion-1.4.pt 38 | wget -O kl.pt https://huggingface.co/Jack000/glid-3-xl-stable/resolve/main/default/kl-1.4.pt 39 | 40 | ``` 41 | 42 | # Generating images 43 | note: best results at 512x512 image size 44 | 45 | ``` 46 | python sample.py --model_path diffusion.pt --batch_size 3 --num_batches 3 --text "a cyberpunk girl with a scifi neuralink device on her head" 47 | 48 | # sample with an init image 49 | python sample.py --init_image picture.jpg --skip_timesteps 20 --model_path diffusion.pt --batch_size 3 --num_batches 3 --text "a cyberpunk girl with a scifi neuralink device on her head" 50 | 51 | # generated images saved to ./output/ 52 | # generated image embeddings saved to ./output_npy/ as npy files 53 | ``` 54 | 55 | # Inpainting 56 | A custom inpainting/outpainting model trained for an additional 100k steps 57 | 58 | ``` 59 | # install PyQt5 if you want to use a gui, otherwise supply a mask file 60 | pip install PyQt5 61 | 62 | # download inpaint model 63 | wget -O inpaint.pt https://huggingface.co/Jack000/glid-3-xl-stable/resolve/main/inpaint/ema_0.9999_100000.pt 64 | 65 | # inpaint with gui 66 | python sample.py --model_path inpaint.pt --edit your-image.png --text "your prompt here" 67 | 68 | # the previously painted mask is saved as mask.png 69 | python sample.py --model_path inpaint.pt --edit your-image.png --text "your prompt here" --mask mask.png 70 | 71 | # partial inpaint by skipping timesteps 72 | python sample.py --model_path inpaint.pt --edit your-image.png --text "your prompt here" --skip_timesteps 20 73 | 74 | # outpaint extends the canvas 75 | # --outpaint options: expand, wider, taller, left, top, right, bottom 76 | python sample.py --model_path inpaint.pt --edit your-image.png --text "your prompt here" --outpaint wider 77 | ``` 78 | 79 | # Generate with classifier guidance 80 | note: best results with --ddim --steps 100 81 | 82 | ``` 83 | # download photo classifier 84 | wget -O class-photo.pt https://huggingface.co/Jack000/glid-3-xl-stable/resolve/main/classifier_photo/model060000.pt 85 | 86 | # download art classifier 87 | wget -O class-art.pt https://huggingface.co/Jack000/glid-3-xl-stable/resolve/main/classifier_art/model110000.pt 88 | 89 | # download anime classifier 90 | wget -O class-anime.pt https://huggingface.co/Jack000/glid-3-xl-stable/resolve/main/classifier_anime/model070000.pt 91 | 92 | # generate 93 | python sample.py --ddim --steps 100 --classifier_scale 100 --classifier class-anime.pt --model_path diffusion.pt --text "anime Elon Musk" 94 | 95 | ``` 96 | 97 | # Upscaling 98 | note: best results at 512x512 input and 1024x1024 output (default settings) 11gb vram required 99 | ``` 100 | # download model 101 | wget -O upscale.pt https://huggingface.co/Jack000/glid-3-xl-stable/resolve/main/super_lg/ema_0.999_040000.pt 102 | 103 | python super_large.py --image img.png --skip_timesteps 1 104 | 105 | ``` 106 | 107 | # Training/Fine tuning 108 | Train with same flags as guided diffusion. Data directory should contain image and text files with the same name (image1.png image1.txt) 109 | 110 | ``` 111 | # minimum 48gb vram to train 112 | # batch sizes are per-gpu, not total 113 | 114 | MODEL_FLAGS="--actual_image_size 512 --lr_warmup_steps 10000 --ema_rate 0.9999 --attention_resolutions 64,32,16 --class_cond False --diffusion_steps 1000 --image_size 64 --learn_sigma False --noise_schedule linear --num_channels 320 --num_heads 8 --num_res_blocks 2 --resblock_updown False --use_fp16 True --use_scale_shift_norm False " 115 | TRAIN_FLAGS="--lr 5e-5 --batch_size 32 --log_interval 10 --save_interval 5000 --kl_model kl.pt --resume_checkpoint diffusion.pt" 116 | export OPENAI_LOGDIR=./logs/ 117 | python scripts/image_train_stable.py --data_dir /path/to/image-and-text-files $MODEL_FLAGS $TRAIN_FLAGS 118 | 119 | # multi-gpu 120 | mpiexec -n N python scripts/image_train_stable.py --data_dir /path/to/image-and-text-files $MODEL_FLAGS $TRAIN_FLAGS 121 | ``` 122 | 123 | ``` 124 | # merge checkpoint back into single .pt (for compatibility with other stable diffusion tools) 125 | 126 | python merge.py sd-v1-4.ckpt ./logs/finetuned-ema-checkpoint.pt 127 | 128 | ``` 129 | 130 | # Train inpainting 131 | 132 | ``` 133 | # example configs for 8x80GB A100 134 | MODEL_FLAGS="--actual_image_size 512 --lr_warmup_steps 10000 --ema_rate 0.9999 --attention_resolutions 64,32,16 --class_cond False --diffusion_steps 1000 --image_size 64 --learn_sigma False --noise_schedule linear --num_channels 320 --num_heads 8 --num_res_blocks 2 --resblock_updown False --use_fp16 True --use_scale_shift_norm False " 135 | TRAIN_FLAGS="--lr 5e-5 --batch_size 32 --log_interval 10 --save_interval 10000 --kl_model kl.pt --resume_checkpoint inpaint.pt" 136 | export OPENAI_LOGDIR=./logs_inpaint/ 137 | mpiexec -n 8 python scripts/image_train_stable_inpaint.py --data_dir /path/to/text/and/images $MODEL_FLAGS $TRAIN_FLAGS 138 | ``` 139 | 140 | # Train classifier 141 | 142 | ``` 143 | # example configs for 4x3090 144 | MODEL_FLAGS="--actual_image_size 512 --weight_decay 0.15 --classifier_attention_resolutions 64,32,16,8 --image_size 64 --classifier_width 128 --classifier_depth 4 --classifier_use_fp16 True " 145 | TRAIN_FLAGS="--lr 2e-5 --batch_size 20 --log_interval 10 --save_interval 10000 --kl_model kl.pt" 146 | export OPENAI_LOGDIR=./logs_classifier/ 147 | 148 | mpiexec -n 4 python scripts/classifier_train_stable.py --good_data_dir /your-images/ --bad_data_dir /laion-images/ $MODEL_FLAGS $TRAIN_FLAGS 149 | ``` 150 | 151 | # Train large upscale model 152 | 153 | ``` 154 | # minimum 80gb vram to train 155 | MODEL_FLAGS="--actual_image_size 1024 --lr_warmup_steps 1000 --ema_rate 0.999 --weight_decay 0.005 --attention_resolutions 64,32 --class_cond False --diffusion_steps 1000 --image_size 1024 --learn_sigma True --noise_schedule linear_openai --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True " 156 | TRAIN_FLAGS="--lr 5e-5 --batch_size 6 --microbatch 3 --log_interval 1 --save_interval 5000 --kl_model kl.pt --resume_checkpoint 256x256_diffusion_uncond.pt" 157 | export OPENAI_LOGDIR=./logs_super/ 158 | mpiexec -n 8 python scripts/image_train_super.py --data_dir /path/to/image-and-text-files $MODEL_FLAGS $TRAIN_FLAGS 159 | 160 | ``` 161 | -------------------------------------------------------------------------------- /datasets/README.md: -------------------------------------------------------------------------------- 1 | # Downloading datasets 2 | 3 | This directory includes instructions and scripts for downloading ImageNet and LSUN bedrooms for use in this codebase. 4 | 5 | ## Class-conditional ImageNet 6 | 7 | For our class-conditional models, we use the official ILSVRC2012 dataset with manual center cropping and downsampling. To obtain this dataset, navigate to [this page on image-net.org](http://www.image-net.org/challenges/LSVRC/2012/downloads) and sign in (or create an account if you do not already have one). Then click on the link reading "Training images (Task 1 & 2)". This is a 138GB tar file containing 1000 sub-tar files, one per class. 8 | 9 | Once the file is downloaded, extract it and look inside. You should see 1000 `.tar` files. You need to extract each of these, which may be impractical to do by hand on your operating system. To automate the process on a Unix-based system, you can `cd` into the directory and run this short shell script: 10 | 11 | ``` 12 | for file in *.tar; do tar xf "$file"; rm "$file"; done 13 | ``` 14 | 15 | This will extract and remove each tar file in turn. 16 | 17 | Once all of the images have been extracted, the resulting directory should be usable as a data directory (the `--data_dir` argument for the training script). The filenames should all start with WNID (class ids) followed by underscores, like `n01440764_2708.JPEG`. Conveniently (but not by accident) this is how the automated data-loader expects to discover class labels. 18 | 19 | ## LSUN bedroom 20 | 21 | To download and pre-process LSUN bedroom, clone [fyu/lsun](https://github.com/fyu/lsun) on GitHub and run their download script `python3 download.py bedroom`. The result will be an "lmdb" database named like `bedroom_train_lmdb`. You can pass this to our [lsun_bedroom.py](lsun_bedroom.py) script like so: 22 | 23 | ``` 24 | python lsun_bedroom.py bedroom_train_lmdb lsun_train_output_dir 25 | ``` 26 | 27 | This creates a directory called `lsun_train_output_dir`. This directory can be passed to the training scripts via the `--data_dir` argument. 28 | -------------------------------------------------------------------------------- /datasets/lsun_bedroom.py: -------------------------------------------------------------------------------- 1 | """ 2 | Convert an LSUN lmdb database into a directory of images. 3 | """ 4 | 5 | import argparse 6 | import io 7 | import os 8 | 9 | from PIL import Image 10 | import lmdb 11 | import numpy as np 12 | 13 | 14 | def read_images(lmdb_path, image_size): 15 | env = lmdb.open(lmdb_path, map_size=1099511627776, max_readers=100, readonly=True) 16 | with env.begin(write=False) as transaction: 17 | cursor = transaction.cursor() 18 | for _, webp_data in cursor: 19 | img = Image.open(io.BytesIO(webp_data)) 20 | width, height = img.size 21 | scale = image_size / min(width, height) 22 | img = img.resize( 23 | (int(round(scale * width)), int(round(scale * height))), 24 | resample=Image.BOX, 25 | ) 26 | arr = np.array(img) 27 | h, w, _ = arr.shape 28 | h_off = (h - image_size) // 2 29 | w_off = (w - image_size) // 2 30 | arr = arr[h_off : h_off + image_size, w_off : w_off + image_size] 31 | yield arr 32 | 33 | 34 | def dump_images(out_dir, images, prefix): 35 | if not os.path.exists(out_dir): 36 | os.mkdir(out_dir) 37 | for i, img in enumerate(images): 38 | Image.fromarray(img).save(os.path.join(out_dir, f"{prefix}_{i:07d}.png")) 39 | 40 | 41 | def main(): 42 | parser = argparse.ArgumentParser() 43 | parser.add_argument("--image-size", help="new image size", type=int, default=256) 44 | parser.add_argument("--prefix", help="class name", type=str, default="bedroom") 45 | parser.add_argument("lmdb_path", help="path to an LSUN lmdb database") 46 | parser.add_argument("out_dir", help="path to output directory") 47 | args = parser.parse_args() 48 | 49 | images = read_images(args.lmdb_path, args.image_size) 50 | dump_images(args.out_dir, images, args.prefix) 51 | 52 | 53 | if __name__ == "__main__": 54 | main() 55 | -------------------------------------------------------------------------------- /encoders/__pycache__/modules.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jack000/glid-3-xl-stable/2b53826e54c042b5e6feb3154d1d8461213d38ca/encoders/__pycache__/modules.cpython-38.pyc -------------------------------------------------------------------------------- /encoders/__pycache__/x_transformer.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jack000/glid-3-xl-stable/2b53826e54c042b5e6feb3154d1d8461213d38ca/encoders/__pycache__/x_transformer.cpython-38.pyc -------------------------------------------------------------------------------- /encoders/modules.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from functools import partial 4 | 5 | from encoders.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 | -------------------------------------------------------------------------------- /evaluations/README.md: -------------------------------------------------------------------------------- 1 | # Evaluations 2 | 3 | To compare different generative models, we use FID, sFID, Precision, Recall, and Inception Score. These metrics can all be calculated using batches of samples, which we store in `.npz` (numpy) files. 4 | 5 | # Download batches 6 | 7 | We provide pre-computed sample batches for the reference datasets, our diffusion models, and several baselines we compare against. These are all stored in `.npz` format. 8 | 9 | Reference dataset batches contain pre-computed statistics over the whole dataset, as well as 10,000 images for computing Precision and Recall. All other batches contain 50,000 images which can be used to compute statistics and Precision/Recall. 10 | 11 | Here are links to download all of the sample and reference batches: 12 | 13 | * LSUN 14 | * LSUN bedroom: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/bedroom/VIRTUAL_lsun_bedroom256.npz) 15 | * [ADM (dropout)](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/bedroom/admnet_dropout_lsun_bedroom.npz) 16 | * [DDPM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/bedroom/ddpm_lsun_bedroom.npz) 17 | * [IDDPM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/bedroom/iddpm_lsun_bedroom.npz) 18 | * [StyleGAN](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/bedroom/stylegan_lsun_bedroom.npz) 19 | * LSUN cat: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/cat/VIRTUAL_lsun_cat256.npz) 20 | * [ADM (dropout)](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/cat/admnet_dropout_lsun_cat.npz) 21 | * [StyleGAN2](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/cat/stylegan2_lsun_cat.npz) 22 | * LSUN horse: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/horse/VIRTUAL_lsun_horse256.npz) 23 | * [ADM (dropout)](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/horse/admnet_dropout_lsun_horse.npz) 24 | * [ADM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/horse/admnet_lsun_horse.npz) 25 | * ImageNet 64x64: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/64/VIRTUAL_imagenet64_labeled.npz) 26 | * [ADM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/64/admnet_imagenet64.npz) 27 | * [IDDPM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/64/iddpm_imagenet64.npz) 28 | * [BigGAN](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/64/biggan_deep_imagenet64.npz) 29 | * ImageNet 128x128: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/128/VIRTUAL_imagenet128_labeled.npz) 30 | * [ADM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/128/admnet_imagenet128.npz) 31 | * [ADM-G](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/128/admnet_guided_imagenet128.npz) 32 | * [ADM-G, 25 steps](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/128/admnet_guided_25step_imagenet128.npz) 33 | * [BigGAN-deep (trunc=1.0)](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/128/biggan_deep_trunc1_imagenet128.npz) 34 | * ImageNet 256x256: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/VIRTUAL_imagenet256_labeled.npz) 35 | * [ADM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/admnet_imagenet256.npz) 36 | * [ADM-G](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/admnet_guided_imagenet256.npz) 37 | * [ADM-G, 25 step](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/admnet_guided_25step_imagenet256.npz) 38 | * [ADM-G + ADM-U](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/admnet_guided_upsampled_imagenet256.npz) 39 | * [ADM-U](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/admnet_upsampled_imagenet256.npz) 40 | * [BigGAN-deep (trunc=1.0)](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/biggan_deep_trunc1_imagenet256.npz) 41 | * ImageNet 512x512: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/VIRTUAL_imagenet512.npz) 42 | * [ADM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/admnet_imagenet512.npz) 43 | * [ADM-G](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/admnet_guided_imagenet512.npz) 44 | * [ADM-G, 25 step](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/admnet_guided_25step_imagenet512.npz) 45 | * [ADM-G + ADM-U](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/admnet_guided_upsampled_imagenet512.npz) 46 | * [ADM-U](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/admnet_upsampled_imagenet512.npz) 47 | * [BigGAN-deep (trunc=1.0)](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/biggan_deep_trunc1_imagenet512.npz) 48 | 49 | # Run evaluations 50 | 51 | First, generate or download a batch of samples and download the corresponding reference batch for the given dataset. For this example, we'll use ImageNet 256x256, so the refernce batch is `VIRTUAL_imagenet256_labeled.npz` and we can use the sample batch `admnet_guided_upsampled_imagenet256.npz`. 52 | 53 | Next, run the `evaluator.py` script. The requirements of this script can be found in [requirements.txt](requirements.txt). Pass two arguments to the script: the reference batch and the sample batch. The script will download the InceptionV3 model used for evaluations into the current working directory (if it is not already present). This file is roughly 100MB. 54 | 55 | The output of the script will look something like this, where the first `...` is a bunch of verbose TensorFlow logging: 56 | 57 | ``` 58 | $ python evaluator.py VIRTUAL_imagenet256_labeled.npz admnet_guided_upsampled_imagenet256.npz 59 | ... 60 | computing reference batch activations... 61 | computing/reading reference batch statistics... 62 | computing sample batch activations... 63 | computing/reading sample batch statistics... 64 | Computing evaluations... 65 | Inception Score: 215.8370361328125 66 | FID: 3.9425574129223264 67 | sFID: 6.140433703346162 68 | Precision: 0.8265 69 | Recall: 0.5309 70 | ``` 71 | -------------------------------------------------------------------------------- /evaluations/requirements.txt: -------------------------------------------------------------------------------- 1 | tensorflow-gpu>=2.0 2 | scipy 3 | requests 4 | tqdm -------------------------------------------------------------------------------- /guided_diffusion.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 2.1 2 | Name: guided-diffusion 3 | Version: 0.0.0 4 | Summary: UNKNOWN 5 | License: UNKNOWN 6 | Platform: UNKNOWN 7 | License-File: LICENSE 8 | 9 | UNKNOWN 10 | 11 | -------------------------------------------------------------------------------- /guided_diffusion.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | LICENSE 2 | README.md 3 | setup.py 4 | guided_diffusion.egg-info/PKG-INFO 5 | guided_diffusion.egg-info/SOURCES.txt 6 | guided_diffusion.egg-info/dependency_links.txt 7 | guided_diffusion.egg-info/requires.txt 8 | guided_diffusion.egg-info/top_level.txt -------------------------------------------------------------------------------- /guided_diffusion.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /guided_diffusion.egg-info/requires.txt: -------------------------------------------------------------------------------- 1 | blobfile>=1.0.5 2 | torch 3 | tqdm 4 | -------------------------------------------------------------------------------- /guided_diffusion.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | guided_diffusion 2 | -------------------------------------------------------------------------------- /guided_diffusion/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Codebase for "Improved Denoising Diffusion Probabilistic Models". 3 | """ 4 | -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jack000/glid-3-xl-stable/2b53826e54c042b5e6feb3154d1d8461213d38ca/guided_diffusion/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/dist_util.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jack000/glid-3-xl-stable/2b53826e54c042b5e6feb3154d1d8461213d38ca/guided_diffusion/__pycache__/dist_util.cpython-38.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/fp16_util.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jack000/glid-3-xl-stable/2b53826e54c042b5e6feb3154d1d8461213d38ca/guided_diffusion/__pycache__/fp16_util.cpython-38.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/gaussian_diffusion.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jack000/glid-3-xl-stable/2b53826e54c042b5e6feb3154d1d8461213d38ca/guided_diffusion/__pycache__/gaussian_diffusion.cpython-38.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/image_datasets_classifier.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jack000/glid-3-xl-stable/2b53826e54c042b5e6feb3154d1d8461213d38ca/guided_diffusion/__pycache__/image_datasets_classifier.cpython-38.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/image_text_datasets.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jack000/glid-3-xl-stable/2b53826e54c042b5e6feb3154d1d8461213d38ca/guided_diffusion/__pycache__/image_text_datasets.cpython-38.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/logger.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jack000/glid-3-xl-stable/2b53826e54c042b5e6feb3154d1d8461213d38ca/guided_diffusion/__pycache__/logger.cpython-38.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/losses.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jack000/glid-3-xl-stable/2b53826e54c042b5e6feb3154d1d8461213d38ca/guided_diffusion/__pycache__/losses.cpython-38.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/nn.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jack000/glid-3-xl-stable/2b53826e54c042b5e6feb3154d1d8461213d38ca/guided_diffusion/__pycache__/nn.cpython-38.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/resample.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jack000/glid-3-xl-stable/2b53826e54c042b5e6feb3154d1d8461213d38ca/guided_diffusion/__pycache__/resample.cpython-38.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/respace.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jack000/glid-3-xl-stable/2b53826e54c042b5e6feb3154d1d8461213d38ca/guided_diffusion/__pycache__/respace.cpython-38.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/script_util.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jack000/glid-3-xl-stable/2b53826e54c042b5e6feb3154d1d8461213d38ca/guided_diffusion/__pycache__/script_util.cpython-38.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/train_util.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jack000/glid-3-xl-stable/2b53826e54c042b5e6feb3154d1d8461213d38ca/guided_diffusion/__pycache__/train_util.cpython-38.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/unet.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Jack000/glid-3-xl-stable/2b53826e54c042b5e6feb3154d1d8461213d38ca/guided_diffusion/__pycache__/unet.cpython-38.pyc -------------------------------------------------------------------------------- /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 | from mpi4py import MPI 11 | import torch as th 12 | import torch.distributed as dist 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 = 8 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 | # th.cuda.set_device(MPI.COMM_WORLD.Get_rank()) 29 | 30 | comm = MPI.COMM_WORLD 31 | backend = "gloo" if not th.cuda.is_available() else "nccl" 32 | 33 | if backend == "gloo": 34 | hostname = "localhost" 35 | else: 36 | hostname = socket.gethostbyname(socket.getfqdn()) 37 | os.environ["MASTER_ADDR"] = comm.bcast(hostname, root=0) 38 | os.environ["RANK"] = str(comm.rank) 39 | os.environ["WORLD_SIZE"] = str(comm.size) 40 | 41 | port = comm.bcast(_find_free_port(), root=0) 42 | os.environ["MASTER_PORT"] = str(port) 43 | dist.init_process_group(backend=backend, init_method="env://") 44 | 45 | 46 | def dev(): 47 | """ 48 | Get the device to use for torch.distributed. 49 | """ 50 | if th.cuda.is_available(): 51 | return th.device(f"cuda") 52 | return th.device("cpu") 53 | 54 | 55 | def load_state_dict(path, **kwargs): 56 | with bf.BlobFile(path, "rb") as f: 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_datasets.py: -------------------------------------------------------------------------------- 1 | import math 2 | import random 3 | 4 | from PIL import Image 5 | import blobfile as bf 6 | from mpi4py import MPI 7 | import numpy as np 8 | from torch.utils.data import DataLoader, Dataset 9 | 10 | def load_data( 11 | *, 12 | data_dir, 13 | batch_size, 14 | image_size, 15 | class_cond=False, 16 | deterministic=False, 17 | random_crop=False, 18 | random_flip=True, 19 | ): 20 | """ 21 | For a dataset, create a generator over (images, kwargs) pairs. 22 | 23 | Each images is an NCHW float tensor, and the kwargs dict contains zero or 24 | more keys, each of which map to a batched Tensor of their own. 25 | The kwargs dict can be used for class labels, in which case the key is "y" 26 | and the values are integer tensors of class labels. 27 | 28 | :param data_dir: a dataset directory. 29 | :param batch_size: the batch size of each returned pair. 30 | :param image_size: the size to which images are resized. 31 | :param class_cond: if True, include a "y" key in returned dicts for class 32 | label. If classes are not available and this is true, an 33 | exception will be raised. 34 | :param deterministic: if True, yield results in a deterministic order. 35 | :param random_crop: if True, randomly crop the images for augmentation. 36 | :param random_flip: if True, randomly flip the images for augmentation. 37 | """ 38 | if not data_dir: 39 | raise ValueError("unspecified data directory") 40 | all_files = _list_image_files_recursively(data_dir) 41 | classes = None 42 | if class_cond: 43 | # Assume classes are the first part of the filename, 44 | # before an underscore. 45 | class_names = [bf.basename(path).split("_")[0] for path in all_files] 46 | sorted_classes = {x: i for i, x in enumerate(sorted(set(class_names)))} 47 | classes = [sorted_classes[x] for x in class_names] 48 | dataset = ImageDataset( 49 | image_size, 50 | all_files, 51 | classes=classes, 52 | shard=MPI.COMM_WORLD.Get_rank(), 53 | num_shards=MPI.COMM_WORLD.Get_size(), 54 | random_crop=random_crop, 55 | random_flip=random_flip, 56 | ) 57 | if deterministic: 58 | loader = DataLoader( 59 | dataset, batch_size=batch_size, shuffle=False, num_workers=1, drop_last=True 60 | ) 61 | else: 62 | loader = DataLoader( 63 | dataset, batch_size=batch_size, shuffle=True, num_workers=1, drop_last=True 64 | ) 65 | while True: 66 | yield from loader 67 | 68 | 69 | def _list_image_files_recursively(data_dir): 70 | results = [] 71 | for entry in sorted(bf.listdir(data_dir)): 72 | full_path = bf.join(data_dir, entry) 73 | ext = entry.split(".")[-1] 74 | if "." in entry and ext.lower() in ["jpg", "jpeg", "png", "gif", "webp"]: 75 | results.append(full_path) 76 | elif bf.isdir(full_path): 77 | results.extend(_list_image_files_recursively(full_path)) 78 | return results 79 | 80 | 81 | class ImageDataset(Dataset): 82 | def __init__( 83 | self, 84 | resolution, 85 | image_paths, 86 | classes=None, 87 | shard=0, 88 | num_shards=1, 89 | random_crop=False, 90 | random_flip=True, 91 | ): 92 | super().__init__() 93 | self.resolution = resolution 94 | self.local_images = image_paths[shard:][::num_shards] 95 | self.local_classes = None if classes is None else classes[shard:][::num_shards] 96 | self.random_crop = random_crop 97 | self.random_flip = random_flip 98 | 99 | def __len__(self): 100 | return len(self.local_images) 101 | 102 | def __getitem__(self, idx): 103 | path = self.local_images[idx] 104 | with bf.BlobFile(path, "rb") as f: 105 | pil_image = Image.open(f) 106 | pil_image.load() 107 | pil_image = pil_image.convert("RGB") 108 | 109 | if self.random_crop: 110 | arr = random_crop_arr(pil_image, self.resolution) 111 | else: 112 | arr = center_crop_arr(pil_image, self.resolution) 113 | 114 | if self.random_flip and random.random() < 0.5: 115 | arr = arr[:, ::-1] 116 | 117 | arr = arr.astype(np.float32) / 127.5 - 1 118 | 119 | out_dict = {} 120 | if self.local_classes is not None: 121 | out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64) 122 | return np.transpose(arr, [2, 0, 1]), out_dict 123 | 124 | 125 | def center_crop_arr(pil_image, image_size): 126 | # We are not on a new enough PIL to support the `reducing_gap` 127 | # argument, which uses BOX downsampling at powers of two first. 128 | # Thus, we do it by hand to improve downsample quality. 129 | while min(*pil_image.size) >= 2 * image_size: 130 | pil_image = pil_image.resize( 131 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX 132 | ) 133 | 134 | scale = image_size / min(*pil_image.size) 135 | pil_image = pil_image.resize( 136 | tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC 137 | ) 138 | 139 | arr = np.array(pil_image) 140 | crop_y = (arr.shape[0] - image_size) // 2 141 | crop_x = (arr.shape[1] - image_size) // 2 142 | return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] 143 | 144 | 145 | def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0): 146 | min_smaller_dim_size = math.ceil(image_size / max_crop_frac) 147 | max_smaller_dim_size = math.ceil(image_size / min_crop_frac) 148 | smaller_dim_size = random.randrange(min_smaller_dim_size, max_smaller_dim_size + 1) 149 | 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 * smaller_dim_size: 154 | pil_image = pil_image.resize( 155 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX 156 | ) 157 | 158 | scale = smaller_dim_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 = random.randrange(arr.shape[0] - image_size + 1) 165 | crop_x = random.randrange(arr.shape[1] - image_size + 1) 166 | return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] 167 | -------------------------------------------------------------------------------- /guided_diffusion/image_datasets_classifier.py: -------------------------------------------------------------------------------- 1 | import math 2 | import random 3 | 4 | from PIL import Image 5 | import blobfile as bf 6 | from mpi4py import MPI 7 | import numpy as np 8 | from torch.utils.data import DataLoader, Dataset 9 | import os 10 | from torchvision import transforms 11 | 12 | def load_data( 13 | *, 14 | good_data_dir, 15 | bad_data_dir, 16 | batch_size, 17 | image_size, 18 | val_split=False, 19 | val_num=0 20 | ): 21 | """ 22 | For a dataset, create a generator over (images, kwargs) pairs. 23 | 24 | Each images is an NCHW float tensor, and the kwargs dict contains zero or 25 | more keys, each of which map to a batched Tensor of their own. 26 | The kwargs dict can be used for class labels, in which case the key is "y" 27 | and the values are integer tensors of class labels. 28 | 29 | :param good_data_dir: images that will be labelled "1" 30 | :param bad_data_dir: images that will be labelled "0" 31 | :param batch_size: the batch size of each returned pair. 32 | :param image_size: the size to which images are resized. 33 | """ 34 | if not good_data_dir or not bad_data_dir: 35 | raise ValueError("unspecified data directory") 36 | 37 | good_files = _list_image_files_recursively(good_data_dir) 38 | bad_files = _list_image_files_recursively(bad_data_dir) 39 | 40 | print('found ', len(good_files),' good files') 41 | print('found ', len(bad_files),' bad files') 42 | 43 | if len(good_files) < len(bad_files): 44 | good_files *= math.floor(len(bad_files)/len(good_files)) 45 | print('duplicated good files to ', len(good_files)) 46 | elif len(bad_files) < len(good_files): 47 | bad_files *= math.floor(len(good_files)/len(bad_files)) 48 | print('duplicated bad files to ', len(bad_files)) 49 | 50 | all_data = [(d,0) for d in bad_files] + [(d,1) for d in good_files] 51 | random.Random(99).shuffle(all_data) 52 | 53 | if val_num > 0: 54 | if val_split: 55 | all_data = all_data[:val_num] 56 | else: 57 | all_data = all_data[val_num:] 58 | 59 | dataset = ImageDataset( 60 | image_size, 61 | all_data, 62 | shard=MPI.COMM_WORLD.Get_rank(), 63 | num_shards=MPI.COMM_WORLD.Get_size(), 64 | ) 65 | loader = DataLoader( 66 | dataset, batch_size=batch_size, shuffle=False, num_workers=1, drop_last=True 67 | ) 68 | while True: 69 | yield from loader 70 | 71 | 72 | def _list_image_files_recursively(data_dir): 73 | results = [] 74 | for entry in sorted(bf.listdir(data_dir)): 75 | full_path = bf.join(data_dir, entry) 76 | entry = entry.split(".") 77 | ext = entry[-1].strip() 78 | filename = entry[0] 79 | if ext and ext.lower() in ["jpg", "jpeg", "png", "gif", "webp"]: 80 | results.append(full_path) 81 | elif bf.isdir(full_path): 82 | results.extend(_list_image_files_recursively(full_path)) 83 | return results 84 | 85 | 86 | class ImageDataset(Dataset): 87 | def __init__( 88 | self, 89 | resolution, 90 | file_paths, 91 | shard=0, 92 | num_shards=1, 93 | ): 94 | super().__init__() 95 | self.resolution = resolution 96 | self.data = file_paths[shard:][::num_shards] 97 | 98 | def __len__(self): 99 | return len(self.data) 100 | 101 | def __getitem__(self, idx): 102 | pdata = self.data[idx] 103 | image_path = pdata[0] 104 | image_class = pdata[1] 105 | 106 | with bf.BlobFile(image_path, "rb") as f: 107 | pil_image = Image.open(f) 108 | pil_image.load() 109 | 110 | pil_image = pil_image.convert("RGB") 111 | 112 | arr = center_crop_arr(pil_image, self.resolution) 113 | 114 | arr = arr.astype(np.float32) / 127.5 - 1 115 | 116 | return np.transpose(arr, [2, 0, 1]), {'y':np.array(image_class, dtype=np.int64)} 117 | 118 | def center_crop_arr(pil_image, image_size): 119 | # We are not on a new enough PIL to support the `reducing_gap` 120 | # argument, which uses BOX downsampling at powers of two first. 121 | # Thus, we do it by hand to improve downsample quality. 122 | while min(*pil_image.size) >= 2 * image_size: 123 | pil_image = pil_image.resize( 124 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX 125 | ) 126 | 127 | scale = image_size / min(*pil_image.size) 128 | pil_image = pil_image.resize( 129 | tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC 130 | ) 131 | 132 | arr = np.array(pil_image) 133 | crop_y = (arr.shape[0] - image_size) // 2 134 | crop_x = (arr.shape[1] - image_size) // 2 135 | return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] 136 | 137 | 138 | def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0): 139 | min_smaller_dim_size = math.ceil(image_size / max_crop_frac) 140 | max_smaller_dim_size = math.ceil(image_size / min_crop_frac) 141 | smaller_dim_size = random.randrange(min_smaller_dim_size, max_smaller_dim_size + 1) 142 | 143 | # We are not on a new enough PIL to support the `reducing_gap` 144 | # argument, which uses BOX downsampling at powers of two first. 145 | # Thus, we do it by hand to improve downsample quality. 146 | while min(*pil_image.size) >= 2 * smaller_dim_size: 147 | pil_image = pil_image.resize( 148 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX 149 | ) 150 | 151 | scale = smaller_dim_size / min(*pil_image.size) 152 | pil_image = pil_image.resize( 153 | tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC 154 | ) 155 | 156 | arr = np.array(pil_image) 157 | crop_y = random.randrange(arr.shape[0] - image_size + 1) 158 | crop_x = random.randrange(arr.shape[1] - image_size + 1) 159 | return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] 160 | -------------------------------------------------------------------------------- /guided_diffusion/image_text_datasets.py: -------------------------------------------------------------------------------- 1 | import math 2 | import random 3 | 4 | from PIL import Image 5 | import blobfile as bf 6 | from mpi4py import MPI 7 | import numpy as np 8 | from torch.utils.data import DataLoader, Dataset 9 | 10 | def load_data( 11 | *, 12 | data_dir, 13 | batch_size, 14 | image_size, 15 | class_cond=False, 16 | deterministic=False, 17 | random_crop=False, 18 | random_flip=True, 19 | ): 20 | """ 21 | For a dataset, create a generator over (images, kwargs) pairs. 22 | 23 | Each images is an NCHW float tensor, and the kwargs dict contains zero or 24 | more keys, each of which map to a batched Tensor of their own. 25 | The kwargs dict can be used for class labels, in which case the key is "y" 26 | and the values are integer tensors of class labels. 27 | 28 | :param data_dir: a dataset directory. 29 | :param batch_size: the batch size of each returned pair. 30 | :param image_size: the size to which images are resized. 31 | :param class_cond: if True, include a "y" key in returned dicts for class 32 | label. If classes are not available and this is true, an 33 | exception will be raised. 34 | :param deterministic: if True, yield results in a deterministic order. 35 | :param random_crop: if True, randomly crop the images for augmentation. 36 | :param random_flip: if True, randomly flip the images for augmentation. 37 | """ 38 | if not data_dir: 39 | raise ValueError("unspecified data directory") 40 | all_files = _list_image_files_recursively(data_dir) 41 | classes = None 42 | #if class_cond: 43 | # Assume classes are the first part of the filename, 44 | # before an underscore. 45 | # class_names = [bf.basename(path).split("_")[0] for path in all_files] 46 | # sorted_classes = {x: i for i, x in enumerate(sorted(set(class_names)))} 47 | # classes = [sorted_classes[x] for x in class_names] 48 | dataset = ImageDataset( 49 | image_size, 50 | all_files, 51 | classes=classes, 52 | shard=MPI.COMM_WORLD.Get_rank(), 53 | num_shards=MPI.COMM_WORLD.Get_size(), 54 | random_crop=random_crop, 55 | random_flip=random_flip, 56 | ) 57 | if deterministic: 58 | loader = DataLoader( 59 | dataset, batch_size=batch_size, shuffle=False, num_workers=1, drop_last=True 60 | ) 61 | else: 62 | loader = DataLoader( 63 | dataset, batch_size=batch_size, shuffle=True, num_workers=1, drop_last=True 64 | ) 65 | while True: 66 | yield from loader 67 | 68 | 69 | def _list_image_files_recursively(data_dir): 70 | results = [] 71 | for entry in sorted(bf.listdir(data_dir)): 72 | full_path = bf.join(data_dir, entry) 73 | entry = entry.split(".") 74 | ext = entry[-1].strip() 75 | filename = entry[0] 76 | if ext and ext.lower() in ["jpg", "jpeg", "png", "gif", "webp"]: 77 | text_path = bf.join(data_dir, filename+'.txt') 78 | if bf.exists(text_path): 79 | results.append((full_path, text_path)) 80 | elif bf.isdir(full_path): 81 | results.extend(_list_image_files_recursively(full_path)) 82 | return results 83 | 84 | 85 | class ImageDataset(Dataset): 86 | def __init__( 87 | self, 88 | resolution, 89 | file_paths, 90 | classes=None, 91 | shard=0, 92 | num_shards=1, 93 | random_crop=False, 94 | random_flip=True, 95 | ): 96 | super().__init__() 97 | self.resolution = resolution 98 | self.local_files = file_paths[shard:][::num_shards] 99 | self.local_classes = None if classes is None else classes[shard:][::num_shards] 100 | self.random_crop = random_crop 101 | self.random_flip = random_flip 102 | 103 | def __len__(self): 104 | return len(self.local_files) 105 | 106 | def __getitem__(self, idx): 107 | path = self.local_files[idx] 108 | with bf.BlobFile(path[0], "rb") as f: 109 | pil_image = Image.open(f) 110 | pil_image.load() 111 | pil_image = pil_image.convert("RGB") 112 | 113 | if self.random_crop: 114 | arr = random_crop_arr(pil_image, self.resolution) 115 | else: 116 | arr = center_crop_arr(pil_image, self.resolution) 117 | 118 | if self.random_flip and random.random() < 0.5: 119 | arr = arr[:, ::-1] 120 | 121 | arr = arr.astype(np.float32) / 127.5 - 1 122 | 123 | out_dict = {} 124 | if self.local_classes is not None: 125 | out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64) 126 | 127 | with bf.BlobFile(path[1], "r") as f: 128 | text = f.read().strip() 129 | 130 | return np.transpose(arr, [2, 0, 1]), out_dict, text 131 | 132 | 133 | def center_crop_arr(pil_image, image_size): 134 | # We are not on a new enough PIL to support the `reducing_gap` 135 | # argument, which uses BOX downsampling at powers of two first. 136 | # Thus, we do it by hand to improve downsample quality. 137 | while min(*pil_image.size) >= 2 * image_size: 138 | pil_image = pil_image.resize( 139 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX 140 | ) 141 | 142 | scale = image_size / min(*pil_image.size) 143 | pil_image = pil_image.resize( 144 | tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC 145 | ) 146 | 147 | arr = np.array(pil_image) 148 | crop_y = (arr.shape[0] - image_size) // 2 149 | crop_x = (arr.shape[1] - image_size) // 2 150 | return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] 151 | 152 | 153 | def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0): 154 | min_smaller_dim_size = math.ceil(image_size / max_crop_frac) 155 | max_smaller_dim_size = math.ceil(image_size / min_crop_frac) 156 | smaller_dim_size = random.randrange(min_smaller_dim_size, max_smaller_dim_size + 1) 157 | 158 | # We are not on a new enough PIL to support the `reducing_gap` 159 | # argument, which uses BOX downsampling at powers of two first. 160 | # Thus, we do it by hand to improve downsample quality. 161 | while min(*pil_image.size) >= 2 * smaller_dim_size: 162 | pil_image = pil_image.resize( 163 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX 164 | ) 165 | 166 | scale = smaller_dim_size / min(*pil_image.size) 167 | pil_image = pil_image.resize( 168 | tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC 169 | ) 170 | 171 | arr = np.array(pil_image) 172 | crop_y = random.randrange(arr.shape[0] - image_size + 1) 173 | crop_x = random.randrange(arr.shape[1] - image_size + 1) 174 | return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] 175 | -------------------------------------------------------------------------------- /guided_diffusion/logger.py: -------------------------------------------------------------------------------- 1 | """ 2 | Logger copied from OpenAI baselines to avoid extra RL-based dependencies: 3 | https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/logger.py 4 | """ 5 | 6 | import os 7 | import sys 8 | import shutil 9 | import os.path as osp 10 | import json 11 | import time 12 | import datetime 13 | import tempfile 14 | import warnings 15 | from collections import defaultdict 16 | from contextlib import contextmanager 17 | 18 | DEBUG = 10 19 | INFO = 20 20 | WARN = 30 21 | ERROR = 40 22 | 23 | DISABLED = 50 24 | 25 | 26 | class KVWriter(object): 27 | def writekvs(self, kvs): 28 | raise NotImplementedError 29 | 30 | 31 | class SeqWriter(object): 32 | def writeseq(self, seq): 33 | raise NotImplementedError 34 | 35 | 36 | class HumanOutputFormat(KVWriter, SeqWriter): 37 | def __init__(self, filename_or_file): 38 | if isinstance(filename_or_file, str): 39 | self.file = open(filename_or_file, "wt") 40 | self.own_file = True 41 | else: 42 | assert hasattr(filename_or_file, "read"), ( 43 | "expected file or str, got %s" % filename_or_file 44 | ) 45 | self.file = filename_or_file 46 | self.own_file = False 47 | 48 | def writekvs(self, kvs): 49 | # Create strings for printing 50 | key2str = {} 51 | for (key, val) in sorted(kvs.items()): 52 | if hasattr(val, "__float__"): 53 | valstr = "%-8.3g" % val 54 | else: 55 | valstr = str(val) 56 | key2str[self._truncate(key)] = self._truncate(valstr) 57 | 58 | # Find max widths 59 | if len(key2str) == 0: 60 | print("WARNING: tried to write empty key-value dict") 61 | return 62 | else: 63 | keywidth = max(map(len, key2str.keys())) 64 | valwidth = max(map(len, key2str.values())) 65 | 66 | # Write out the data 67 | dashes = "-" * (keywidth + valwidth + 7) 68 | lines = [dashes] 69 | for (key, val) in sorted(key2str.items(), key=lambda kv: kv[0].lower()): 70 | lines.append( 71 | "| %s%s | %s%s |" 72 | % (key, " " * (keywidth - len(key)), val, " " * (valwidth - len(val))) 73 | ) 74 | lines.append(dashes) 75 | self.file.write("\n".join(lines) + "\n") 76 | 77 | # Flush the output to the file 78 | self.file.flush() 79 | 80 | def _truncate(self, s): 81 | maxlen = 30 82 | return s[: maxlen - 3] + "..." if len(s) > maxlen else s 83 | 84 | def writeseq(self, seq): 85 | seq = list(seq) 86 | for (i, elem) in enumerate(seq): 87 | self.file.write(elem) 88 | if i < len(seq) - 1: # add space unless this is the last one 89 | self.file.write(" ") 90 | self.file.write("\n") 91 | self.file.flush() 92 | 93 | def close(self): 94 | if self.own_file: 95 | self.file.close() 96 | 97 | 98 | class JSONOutputFormat(KVWriter): 99 | def __init__(self, filename): 100 | self.file = open(filename, "wt") 101 | 102 | def writekvs(self, kvs): 103 | for k, v in sorted(kvs.items()): 104 | if hasattr(v, "dtype"): 105 | kvs[k] = float(v) 106 | self.file.write(json.dumps(kvs) + "\n") 107 | self.file.flush() 108 | 109 | def close(self): 110 | self.file.close() 111 | 112 | 113 | class CSVOutputFormat(KVWriter): 114 | def __init__(self, filename): 115 | self.file = open(filename, "w+t") 116 | self.keys = [] 117 | self.sep = "," 118 | 119 | def writekvs(self, kvs): 120 | # Add our current row to the history 121 | extra_keys = list(kvs.keys() - self.keys) 122 | extra_keys.sort() 123 | if extra_keys: 124 | self.keys.extend(extra_keys) 125 | self.file.seek(0) 126 | lines = self.file.readlines() 127 | self.file.seek(0) 128 | for (i, k) in enumerate(self.keys): 129 | if i > 0: 130 | self.file.write(",") 131 | self.file.write(k) 132 | self.file.write("\n") 133 | for line in lines[1:]: 134 | self.file.write(line[:-1]) 135 | self.file.write(self.sep * len(extra_keys)) 136 | self.file.write("\n") 137 | for (i, k) in enumerate(self.keys): 138 | if i > 0: 139 | self.file.write(",") 140 | v = kvs.get(k) 141 | if v is not None: 142 | self.file.write(str(v)) 143 | self.file.write("\n") 144 | self.file.flush() 145 | 146 | def close(self): 147 | self.file.close() 148 | 149 | 150 | class TensorBoardOutputFormat(KVWriter): 151 | """ 152 | Dumps key/value pairs into TensorBoard's numeric format. 153 | """ 154 | 155 | def __init__(self, dir): 156 | os.makedirs(dir, exist_ok=True) 157 | self.dir = dir 158 | self.step = 1 159 | prefix = "events" 160 | path = osp.join(osp.abspath(dir), prefix) 161 | import tensorflow as tf 162 | from tensorflow.python import pywrap_tensorflow 163 | from tensorflow.core.util import event_pb2 164 | from tensorflow.python.util import compat 165 | 166 | self.tf = tf 167 | self.event_pb2 = event_pb2 168 | self.pywrap_tensorflow = pywrap_tensorflow 169 | self.writer = pywrap_tensorflow.EventsWriter(compat.as_bytes(path)) 170 | 171 | def writekvs(self, kvs): 172 | def summary_val(k, v): 173 | kwargs = {"tag": k, "simple_value": float(v)} 174 | return self.tf.Summary.Value(**kwargs) 175 | 176 | summary = self.tf.Summary(value=[summary_val(k, v) for k, v in kvs.items()]) 177 | event = self.event_pb2.Event(wall_time=time.time(), summary=summary) 178 | event.step = ( 179 | self.step 180 | ) # is there any reason why you'd want to specify the step? 181 | self.writer.WriteEvent(event) 182 | self.writer.Flush() 183 | self.step += 1 184 | 185 | def close(self): 186 | if self.writer: 187 | self.writer.Close() 188 | self.writer = None 189 | 190 | 191 | def make_output_format(format, ev_dir, log_suffix=""): 192 | os.makedirs(ev_dir, exist_ok=True) 193 | if format == "stdout": 194 | return HumanOutputFormat(sys.stdout) 195 | elif format == "log": 196 | return HumanOutputFormat(osp.join(ev_dir, "log%s.txt" % log_suffix)) 197 | elif format == "json": 198 | return JSONOutputFormat(osp.join(ev_dir, "progress%s.json" % log_suffix)) 199 | elif format == "csv": 200 | return CSVOutputFormat(osp.join(ev_dir, "progress%s.csv" % log_suffix)) 201 | elif format == "tensorboard": 202 | return TensorBoardOutputFormat(osp.join(ev_dir, "tb%s" % log_suffix)) 203 | else: 204 | raise ValueError("Unknown format specified: %s" % (format,)) 205 | 206 | 207 | # ================================================================ 208 | # API 209 | # ================================================================ 210 | 211 | 212 | def logkv(key, val): 213 | """ 214 | Log a value of some diagnostic 215 | Call this once for each diagnostic quantity, each iteration 216 | If called many times, last value will be used. 217 | """ 218 | get_current().logkv(key, val) 219 | 220 | 221 | def logkv_mean(key, val): 222 | """ 223 | The same as logkv(), but if called many times, values averaged. 224 | """ 225 | get_current().logkv_mean(key, val) 226 | 227 | 228 | def logkvs(d): 229 | """ 230 | Log a dictionary of key-value pairs 231 | """ 232 | for (k, v) in d.items(): 233 | logkv(k, v) 234 | 235 | 236 | def dumpkvs(): 237 | """ 238 | Write all of the diagnostics from the current iteration 239 | """ 240 | return get_current().dumpkvs() 241 | 242 | 243 | def getkvs(): 244 | return get_current().name2val 245 | 246 | 247 | def log(*args, level=INFO): 248 | """ 249 | Write the sequence of args, with no separators, to the console and output files (if you've configured an output file). 250 | """ 251 | get_current().log(*args, level=level) 252 | 253 | 254 | def debug(*args): 255 | log(*args, level=DEBUG) 256 | 257 | 258 | def info(*args): 259 | log(*args, level=INFO) 260 | 261 | 262 | def warn(*args): 263 | log(*args, level=WARN) 264 | 265 | 266 | def error(*args): 267 | log(*args, level=ERROR) 268 | 269 | 270 | def set_level(level): 271 | """ 272 | Set logging threshold on current logger. 273 | """ 274 | get_current().set_level(level) 275 | 276 | 277 | def set_comm(comm): 278 | get_current().set_comm(comm) 279 | 280 | 281 | def get_dir(): 282 | """ 283 | Get directory that log files are being written to. 284 | will be None if there is no output directory (i.e., if you didn't call start) 285 | """ 286 | return get_current().get_dir() 287 | 288 | 289 | record_tabular = logkv 290 | dump_tabular = dumpkvs 291 | 292 | 293 | @contextmanager 294 | def profile_kv(scopename): 295 | logkey = "wait_" + scopename 296 | tstart = time.time() 297 | try: 298 | yield 299 | finally: 300 | get_current().name2val[logkey] += time.time() - tstart 301 | 302 | 303 | def profile(n): 304 | """ 305 | Usage: 306 | @profile("my_func") 307 | def my_func(): code 308 | """ 309 | 310 | def decorator_with_name(func): 311 | def func_wrapper(*args, **kwargs): 312 | with profile_kv(n): 313 | return func(*args, **kwargs) 314 | 315 | return func_wrapper 316 | 317 | return decorator_with_name 318 | 319 | 320 | # ================================================================ 321 | # Backend 322 | # ================================================================ 323 | 324 | 325 | def get_current(): 326 | if Logger.CURRENT is None: 327 | _configure_default_logger() 328 | 329 | return Logger.CURRENT 330 | 331 | 332 | class Logger(object): 333 | DEFAULT = None # A logger with no output files. (See right below class definition) 334 | # So that you can still log to the terminal without setting up any output files 335 | CURRENT = None # Current logger being used by the free functions above 336 | 337 | def __init__(self, dir, output_formats, comm=None): 338 | self.name2val = defaultdict(float) # values this iteration 339 | self.name2cnt = defaultdict(int) 340 | self.level = INFO 341 | self.dir = dir 342 | self.output_formats = output_formats 343 | self.comm = comm 344 | 345 | # Logging API, forwarded 346 | # ---------------------------------------- 347 | def logkv(self, key, val): 348 | self.name2val[key] = val 349 | 350 | def logkv_mean(self, key, val): 351 | oldval, cnt = self.name2val[key], self.name2cnt[key] 352 | self.name2val[key] = oldval * cnt / (cnt + 1) + val / (cnt + 1) 353 | self.name2cnt[key] = cnt + 1 354 | 355 | def dumpkvs(self): 356 | if self.comm is None: 357 | d = self.name2val 358 | else: 359 | d = mpi_weighted_mean( 360 | self.comm, 361 | { 362 | name: (val, self.name2cnt.get(name, 1)) 363 | for (name, val) in self.name2val.items() 364 | }, 365 | ) 366 | if self.comm.rank != 0: 367 | d["dummy"] = 1 # so we don't get a warning about empty dict 368 | out = d.copy() # Return the dict for unit testing purposes 369 | for fmt in self.output_formats: 370 | if isinstance(fmt, KVWriter): 371 | fmt.writekvs(d) 372 | self.name2val.clear() 373 | self.name2cnt.clear() 374 | return out 375 | 376 | def log(self, *args, level=INFO): 377 | if self.level <= level: 378 | self._do_log(args) 379 | 380 | # Configuration 381 | # ---------------------------------------- 382 | def set_level(self, level): 383 | self.level = level 384 | 385 | def set_comm(self, comm): 386 | self.comm = comm 387 | 388 | def get_dir(self): 389 | return self.dir 390 | 391 | def close(self): 392 | for fmt in self.output_formats: 393 | fmt.close() 394 | 395 | # Misc 396 | # ---------------------------------------- 397 | def _do_log(self, args): 398 | for fmt in self.output_formats: 399 | if isinstance(fmt, SeqWriter): 400 | fmt.writeseq(map(str, args)) 401 | 402 | 403 | def get_rank_without_mpi_import(): 404 | # check environment variables here instead of importing mpi4py 405 | # to avoid calling MPI_Init() when this module is imported 406 | for varname in ["PMI_RANK", "OMPI_COMM_WORLD_RANK"]: 407 | if varname in os.environ: 408 | return int(os.environ[varname]) 409 | return 0 410 | 411 | 412 | def mpi_weighted_mean(comm, local_name2valcount): 413 | """ 414 | Copied from: https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/common/mpi_util.py#L110 415 | Perform a weighted average over dicts that are each on a different node 416 | Input: local_name2valcount: dict mapping key -> (value, count) 417 | Returns: key -> mean 418 | """ 419 | all_name2valcount = comm.gather(local_name2valcount) 420 | if comm.rank == 0: 421 | name2sum = defaultdict(float) 422 | name2count = defaultdict(float) 423 | for n2vc in all_name2valcount: 424 | for (name, (val, count)) in n2vc.items(): 425 | try: 426 | val = float(val) 427 | except ValueError: 428 | if comm.rank == 0: 429 | warnings.warn( 430 | "WARNING: tried to compute mean on non-float {}={}".format( 431 | name, val 432 | ) 433 | ) 434 | else: 435 | name2sum[name] += val * count 436 | name2count[name] += count 437 | return {name: name2sum[name] / name2count[name] for name in name2sum} 438 | else: 439 | return {} 440 | 441 | 442 | def configure(dir=None, format_strs=None, comm=None, log_suffix=""): 443 | """ 444 | If comm is provided, average all numerical stats across that comm 445 | """ 446 | if dir is None: 447 | dir = os.getenv("OPENAI_LOGDIR") 448 | if dir is None: 449 | dir = osp.join( 450 | tempfile.gettempdir(), 451 | datetime.datetime.now().strftime("openai-%Y-%m-%d-%H-%M-%S-%f"), 452 | ) 453 | assert isinstance(dir, str) 454 | dir = os.path.expanduser(dir) 455 | os.makedirs(os.path.expanduser(dir), exist_ok=True) 456 | 457 | rank = get_rank_without_mpi_import() 458 | if rank > 0: 459 | log_suffix = log_suffix + "-rank%03i" % rank 460 | 461 | if format_strs is None: 462 | if rank == 0: 463 | format_strs = os.getenv("OPENAI_LOG_FORMAT", "stdout,log,csv").split(",") 464 | else: 465 | format_strs = os.getenv("OPENAI_LOG_FORMAT_MPI", "log").split(",") 466 | format_strs = filter(None, format_strs) 467 | output_formats = [make_output_format(f, dir, log_suffix) for f in format_strs] 468 | 469 | Logger.CURRENT = Logger(dir=dir, output_formats=output_formats, comm=comm) 470 | if output_formats: 471 | log("Logging to %s" % dir) 472 | 473 | 474 | def _configure_default_logger(): 475 | configure() 476 | Logger.DEFAULT = Logger.CURRENT 477 | 478 | 479 | def reset(): 480 | if Logger.CURRENT is not Logger.DEFAULT: 481 | Logger.CURRENT.close() 482 | Logger.CURRENT = Logger.DEFAULT 483 | log("Reset logger") 484 | 485 | 486 | @contextmanager 487 | def scoped_configure(dir=None, format_strs=None, comm=None): 488 | prevlogger = Logger.CURRENT 489 | configure(dir=dir, format_strs=format_strs, comm=comm) 490 | try: 491 | yield 492 | finally: 493 | Logger.CURRENT.close() 494 | Logger.CURRENT = prevlogger 495 | 496 | -------------------------------------------------------------------------------- /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 | 9 | import torch as th 10 | 11 | 12 | def normal_kl(mean1, logvar1, mean2, logvar2): 13 | """ 14 | Compute the KL divergence between two gaussians. 15 | 16 | Shapes are automatically broadcasted, so batches can be compared to 17 | scalars, among other use cases. 18 | """ 19 | tensor = None 20 | for obj in (mean1, logvar1, mean2, logvar2): 21 | if isinstance(obj, th.Tensor): 22 | tensor = obj 23 | break 24 | assert tensor is not None, "at least one argument must be a Tensor" 25 | 26 | # Force variances to be Tensors. Broadcasting helps convert scalars to 27 | # Tensors, but it does not work for th.exp(). 28 | logvar1, logvar2 = [ 29 | x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor) 30 | for x in (logvar1, logvar2) 31 | ] 32 | 33 | return 0.5 * ( 34 | -1.0 35 | + logvar2 36 | - logvar1 37 | + th.exp(logvar1 - logvar2) 38 | + ((mean1 - mean2) ** 2) * th.exp(-logvar2) 39 | ) 40 | 41 | 42 | def approx_standard_normal_cdf(x): 43 | """ 44 | A fast approximation of the cumulative distribution function of the 45 | standard normal. 46 | """ 47 | return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3)))) 48 | 49 | 50 | def discretized_gaussian_log_likelihood(x, *, means, log_scales): 51 | """ 52 | Compute the log-likelihood of a Gaussian distribution discretizing to a 53 | given image. 54 | 55 | :param x: the target images. It is assumed that this was uint8 values, 56 | rescaled to the range [-1, 1]. 57 | :param means: the Gaussian mean Tensor. 58 | :param log_scales: the Gaussian log stddev Tensor. 59 | :return: a tensor like x of log probabilities (in nats). 60 | """ 61 | assert x.shape == means.shape == log_scales.shape 62 | centered_x = x - means 63 | inv_stdv = th.exp(-log_scales) 64 | plus_in = inv_stdv * (centered_x + 1.0 / 255.0) 65 | cdf_plus = approx_standard_normal_cdf(plus_in) 66 | min_in = inv_stdv * (centered_x - 1.0 / 255.0) 67 | cdf_min = approx_standard_normal_cdf(min_in) 68 | log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12)) 69 | log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12)) 70 | cdf_delta = cdf_plus - cdf_min 71 | log_probs = th.where( 72 | x < -0.999, 73 | log_cdf_plus, 74 | th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))), 75 | ) 76 | assert log_probs.shape == x.shape 77 | return log_probs 78 | -------------------------------------------------------------------------------- /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 | class GroupNorm32(nn.GroupNorm): 12 | def __init__(self, num_groups, num_channels, swish, eps=1e-5): 13 | super().__init__(num_groups=num_groups, num_channels=num_channels, eps=eps) 14 | self.swish = swish 15 | 16 | def forward(self, x): 17 | y = super().forward(x.float()).to(x.dtype) 18 | if self.swish == 1.0: 19 | y = F.silu(y) 20 | elif self.swish: 21 | y = y * F.sigmoid(y * float(self.swish)) 22 | return y 23 | 24 | def conv_nd(dims, *args, **kwargs): 25 | """ 26 | Create a 1D, 2D, or 3D convolution module. 27 | """ 28 | if dims == 1: 29 | return nn.Conv1d(*args, **kwargs) 30 | elif dims == 2: 31 | return nn.Conv2d(*args, **kwargs) 32 | elif dims == 3: 33 | return nn.Conv3d(*args, **kwargs) 34 | raise ValueError(f"unsupported dimensions: {dims}") 35 | 36 | 37 | def linear(*args, **kwargs): 38 | """ 39 | Create a linear module. 40 | """ 41 | return nn.Linear(*args, **kwargs) 42 | 43 | 44 | def avg_pool_nd(dims, *args, **kwargs): 45 | """ 46 | Create a 1D, 2D, or 3D average pooling module. 47 | """ 48 | if dims == 1: 49 | return nn.AvgPool1d(*args, **kwargs) 50 | elif dims == 2: 51 | return nn.AvgPool2d(*args, **kwargs) 52 | elif dims == 3: 53 | return nn.AvgPool3d(*args, **kwargs) 54 | raise ValueError(f"unsupported dimensions: {dims}") 55 | 56 | 57 | def update_ema(target_params, source_params, rate=0.99): 58 | """ 59 | Update target parameters to be closer to those of source parameters using 60 | an exponential moving average. 61 | 62 | :param target_params: the target parameter sequence. 63 | :param source_params: the source parameter sequence. 64 | :param rate: the EMA rate (closer to 1 means slower). 65 | """ 66 | for targ, src in zip(target_params, source_params): 67 | targ.detach().mul_(rate).add_(src, alpha=1 - rate) 68 | 69 | 70 | def zero_module(module): 71 | """ 72 | Zero out the parameters of a module and return it. 73 | """ 74 | for p in module.parameters(): 75 | p.detach().zero_() 76 | return module 77 | 78 | 79 | def scale_module(module, scale): 80 | """ 81 | Scale the parameters of a module and return it. 82 | """ 83 | for p in module.parameters(): 84 | p.detach().mul_(scale) 85 | return module 86 | 87 | 88 | def mean_flat(tensor): 89 | """ 90 | Take the mean over all non-batch dimensions. 91 | """ 92 | return tensor.mean(dim=list(range(1, len(tensor.shape)))) 93 | 94 | 95 | def normalization(channels, swish=0.0): 96 | """ 97 | Make a standard normalization layer, with an optional swish activation. 98 | 99 | :param channels: number of input channels. 100 | :return: an nn.Module for normalization. 101 | """ 102 | return GroupNorm32(num_channels=channels, num_groups=32, swish=swish) 103 | 104 | 105 | #def timestep_embedding(timesteps, dim, max_period=10000): 106 | # """ 107 | # Create sinusoidal timestep embeddings. 108 | 109 | # :param timesteps: a 1-D Tensor of N indices, one per batch element. 110 | # These may be fractional. 111 | # :param dim: the dimension of the output. 112 | # :param max_period: controls the minimum frequency of the embeddings. 113 | # :return: an [N x dim] Tensor of positional embeddings. 114 | # """ 115 | # half = dim // 2 116 | # freqs = th.exp( 117 | # -math.log(max_period) * th.arange(start=0, end=half, dtype=th.float32) / half 118 | # ).to(device=timesteps.device) 119 | # args = timesteps[:, None].float() * freqs[None] 120 | # embedding = th.cat([th.cos(args), th.sin(args)], dim=-1) 121 | # if dim % 2: 122 | # embedding = th.cat([embedding, th.zeros_like(embedding[:, :1])], dim=-1) 123 | # return embedding 124 | 125 | def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False): 126 | """ 127 | Create sinusoidal timestep embeddings. 128 | :param timesteps: a 1-D Tensor of N indices, one per batch element. 129 | These may be fractional. 130 | :param dim: the dimension of the output. 131 | :param max_period: controls the minimum frequency of the embeddings. 132 | :return: an [N x dim] Tensor of positional embeddings. 133 | """ 134 | if not repeat_only: 135 | half = dim // 2 136 | freqs = th.exp( 137 | -math.log(max_period) * th.arange(start=0, end=half, dtype=th.float32) / half 138 | ).to(device=timesteps.device) 139 | args = timesteps[:, None].float() * freqs[None] 140 | embedding = th.cat([th.cos(args), th.sin(args)], dim=-1) 141 | if dim % 2: 142 | embedding = th.cat([embedding, th.zeros_like(embedding[:, :1])], dim=-1) 143 | else: 144 | embedding = repeat(timesteps, 'b -> b d', d=dim) 145 | return embedding 146 | 147 | 148 | 149 | def checkpoint(func, inputs, params, flag): 150 | """ 151 | Evaluate a function without caching intermediate activations, allowing for 152 | reduced memory at the expense of extra compute in the backward pass. 153 | 154 | :param func: the function to evaluate. 155 | :param inputs: the argument sequence to pass to `func`. 156 | :param params: a sequence of parameters `func` depends on but does not 157 | explicitly take as arguments. 158 | :param flag: if False, disable gradient checkpointing. 159 | """ 160 | if flag: 161 | args = tuple(inputs) + tuple(params) 162 | return CheckpointFunction.apply(func, len(inputs), *args) 163 | else: 164 | return func(*inputs) 165 | 166 | 167 | class CheckpointFunction(th.autograd.Function): 168 | @staticmethod 169 | def forward(ctx, run_function, length, *args): 170 | ctx.run_function = run_function 171 | ctx.input_tensors = list(args[:length]) 172 | ctx.input_params = list(args[length:]) 173 | with th.no_grad(): 174 | output_tensors = ctx.run_function(*ctx.input_tensors) 175 | return output_tensors 176 | 177 | @staticmethod 178 | def backward(ctx, *output_grads): 179 | ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors] 180 | with th.enable_grad(): 181 | # Fixes a bug where the first op in run_function modifies the 182 | # Tensor storage in place, which is not allowed for detach()'d 183 | # Tensors. 184 | shallow_copies = [x.view_as(x) for x in ctx.input_tensors] 185 | output_tensors = ctx.run_function(*shallow_copies) 186 | input_grads = th.autograd.grad( 187 | output_tensors, 188 | ctx.input_tensors + ctx.input_params, 189 | output_grads, 190 | allow_unused=True, 191 | ) 192 | del ctx.input_tensors 193 | del ctx.input_params 194 | del output_tensors 195 | return (None, None) + input_grads 196 | -------------------------------------------------------------------------------- /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, EncoderUNetModel 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=True, 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=True, 59 | use_scale_shift_norm=True, 60 | resblock_updown=False, 61 | use_fp16=False, 62 | 63 | use_spatial_transformer=True, 64 | context_dim=768, 65 | 66 | clip_embed_dim=None, 67 | image_condition=False, 68 | super_res_condition=False 69 | ) 70 | res.update(diffusion_defaults()) 71 | return res 72 | 73 | 74 | def classifier_and_diffusion_defaults(): 75 | res = classifier_defaults() 76 | res.update(diffusion_defaults()) 77 | return res 78 | 79 | 80 | def create_model_and_diffusion( 81 | image_size, 82 | class_cond, 83 | learn_sigma, 84 | num_channels, 85 | num_res_blocks, 86 | channel_mult, 87 | num_heads, 88 | num_head_channels, 89 | num_heads_upsample, 90 | attention_resolutions, 91 | dropout, 92 | diffusion_steps, 93 | noise_schedule, 94 | timestep_respacing, 95 | use_kl, 96 | predict_xstart, 97 | rescale_timesteps, 98 | rescale_learned_sigmas, 99 | use_checkpoint, 100 | use_scale_shift_norm, 101 | resblock_updown, 102 | use_fp16, 103 | use_spatial_transformer, 104 | context_dim, 105 | clip_embed_dim, 106 | image_condition, 107 | super_res_condition, 108 | ): 109 | model = create_model( 110 | image_size, 111 | num_channels, 112 | num_res_blocks, 113 | channel_mult=channel_mult, 114 | learn_sigma=learn_sigma, 115 | class_cond=class_cond, 116 | use_checkpoint=use_checkpoint, 117 | attention_resolutions=attention_resolutions, 118 | num_heads=num_heads, 119 | num_head_channels=num_head_channels, 120 | num_heads_upsample=num_heads_upsample, 121 | use_scale_shift_norm=use_scale_shift_norm, 122 | dropout=dropout, 123 | resblock_updown=resblock_updown, 124 | use_fp16=use_fp16, 125 | use_spatial_transformer=use_spatial_transformer, 126 | context_dim=context_dim, 127 | clip_embed_dim=clip_embed_dim, 128 | image_condition=image_condition, 129 | super_res_condition=super_res_condition, 130 | ) 131 | diffusion = create_gaussian_diffusion( 132 | steps=diffusion_steps, 133 | learn_sigma=learn_sigma, 134 | noise_schedule=noise_schedule, 135 | use_kl=use_kl, 136 | predict_xstart=predict_xstart, 137 | rescale_timesteps=rescale_timesteps, 138 | rescale_learned_sigmas=rescale_learned_sigmas, 139 | timestep_respacing=timestep_respacing, 140 | ) 141 | return model, diffusion 142 | 143 | 144 | def create_model( 145 | image_size, 146 | num_channels, 147 | num_res_blocks, 148 | channel_mult="", 149 | learn_sigma=False, 150 | class_cond=False, 151 | use_checkpoint=True, 152 | attention_resolutions="16", 153 | num_heads=1, 154 | num_head_channels=-1, 155 | num_heads_upsample=-1, 156 | use_scale_shift_norm=False, 157 | dropout=0, 158 | resblock_updown=False, 159 | use_fp16=False, 160 | use_spatial_transformer=True, 161 | context_dim=768, 162 | clip_embed_dim=None, 163 | image_condition=False, 164 | super_res_condition=False 165 | ): 166 | if channel_mult == "": 167 | if image_size == 1024: 168 | channel_mult = (1, 1, 2, 2, 4, 4) 169 | elif image_size == 512: 170 | channel_mult = (0.5, 1, 1, 2, 2, 4, 4) 171 | elif image_size == 256: 172 | channel_mult = (1, 1, 2, 2, 4, 4) 173 | elif image_size == 128: 174 | channel_mult = (1, 1, 2, 3, 4) 175 | elif image_size == 64: 176 | channel_mult = (1, 2, 4, 4) 177 | elif image_size == 32: 178 | channel_mult = (1, 2, 4, 4) 179 | else: 180 | raise ValueError(f"unsupported image size: {image_size}") 181 | else: 182 | channel_mult = tuple(float(ch_mult) for ch_mult in channel_mult.split(",")) 183 | 184 | attention_ds = [] 185 | for res in attention_resolutions.split(","): 186 | attention_ds.append(image_size // int(res)) 187 | 188 | if image_size < 512: 189 | in_channels = 4 190 | out_channels = 4 191 | if image_condition: 192 | in_channels = 8 193 | else: 194 | in_channels = 3 195 | out_channels = 3 196 | if learn_sigma: 197 | out_channels = 6 198 | if image_condition: 199 | in_channels = 6 200 | 201 | return UNetModel( 202 | image_size=image_size, 203 | in_channels=in_channels, 204 | model_channels=num_channels, 205 | out_channels=out_channels, 206 | num_res_blocks=num_res_blocks, 207 | attention_resolutions=tuple(attention_ds), 208 | dropout=dropout, 209 | channel_mult=channel_mult, 210 | num_classes=(NUM_CLASSES if class_cond else None), 211 | use_checkpoint=use_checkpoint, 212 | use_fp16=use_fp16, 213 | num_heads=num_heads, 214 | num_head_channels=num_head_channels, 215 | num_heads_upsample=num_heads_upsample, 216 | use_scale_shift_norm=use_scale_shift_norm, 217 | resblock_updown=resblock_updown, 218 | use_spatial_transformer=use_spatial_transformer, 219 | context_dim=context_dim, 220 | clip_embed_dim=clip_embed_dim, 221 | image_condition=image_condition, 222 | super_res_condition=super_res_condition, 223 | ) 224 | 225 | 226 | def create_classifier_and_diffusion( 227 | image_size, 228 | classifier_use_fp16, 229 | classifier_width, 230 | classifier_depth, 231 | classifier_attention_resolutions, 232 | classifier_use_scale_shift_norm, 233 | classifier_resblock_updown, 234 | classifier_pool, 235 | learn_sigma, 236 | diffusion_steps, 237 | noise_schedule, 238 | timestep_respacing, 239 | use_kl, 240 | predict_xstart, 241 | rescale_timesteps, 242 | rescale_learned_sigmas, 243 | ): 244 | classifier = create_classifier( 245 | image_size, 246 | classifier_use_fp16, 247 | classifier_width, 248 | classifier_depth, 249 | classifier_attention_resolutions, 250 | classifier_use_scale_shift_norm, 251 | classifier_resblock_updown, 252 | classifier_pool, 253 | ) 254 | diffusion = create_gaussian_diffusion( 255 | steps=diffusion_steps, 256 | learn_sigma=learn_sigma, 257 | noise_schedule=noise_schedule, 258 | use_kl=use_kl, 259 | predict_xstart=predict_xstart, 260 | rescale_timesteps=rescale_timesteps, 261 | rescale_learned_sigmas=rescale_learned_sigmas, 262 | timestep_respacing=timestep_respacing, 263 | ) 264 | return classifier, diffusion 265 | 266 | 267 | def create_classifier( 268 | image_size, 269 | classifier_use_fp16, 270 | classifier_width, 271 | classifier_depth, 272 | classifier_attention_resolutions, 273 | classifier_use_scale_shift_norm, 274 | classifier_resblock_updown, 275 | classifier_pool, 276 | ): 277 | if image_size == 512: 278 | channel_mult = (0.5, 1, 1, 2, 2, 4, 4) 279 | elif image_size == 256: 280 | channel_mult = (1, 1, 2, 2, 4, 4) 281 | elif image_size == 128: 282 | channel_mult = (1, 1, 2, 3, 4) 283 | elif image_size == 64: 284 | channel_mult = (1, 2, 4, 4) 285 | else: 286 | raise ValueError(f"unsupported image size: {image_size}") 287 | 288 | attention_ds = [] 289 | for res in classifier_attention_resolutions.split(","): 290 | attention_ds.append(image_size // int(res)) 291 | 292 | if image_size < 512: 293 | in_channels = 4 294 | else: 295 | in_channels = 3 296 | 297 | return EncoderUNetModel( 298 | image_size=image_size, 299 | in_channels=in_channels, 300 | model_channels=classifier_width, 301 | out_channels=1000, 302 | num_res_blocks=classifier_depth, 303 | attention_resolutions=tuple(attention_ds), 304 | channel_mult=channel_mult, 305 | use_fp16=classifier_use_fp16, 306 | num_head_channels=64, 307 | use_scale_shift_norm=classifier_use_scale_shift_norm, 308 | resblock_updown=classifier_resblock_updown, 309 | pool=classifier_pool, 310 | ) 311 | 312 | 313 | def sr_model_and_diffusion_defaults(): 314 | res = model_and_diffusion_defaults() 315 | res["large_size"] = 256 316 | res["small_size"] = 64 317 | arg_names = inspect.getfullargspec(sr_create_model_and_diffusion)[0] 318 | for k in res.copy().keys(): 319 | if k not in arg_names: 320 | del res[k] 321 | return res 322 | 323 | 324 | def sr_create_model_and_diffusion( 325 | large_size, 326 | small_size, 327 | class_cond, 328 | learn_sigma, 329 | num_channels, 330 | num_res_blocks, 331 | num_heads, 332 | num_head_channels, 333 | num_heads_upsample, 334 | attention_resolutions, 335 | dropout, 336 | diffusion_steps, 337 | noise_schedule, 338 | timestep_respacing, 339 | use_kl, 340 | predict_xstart, 341 | rescale_timesteps, 342 | rescale_learned_sigmas, 343 | use_checkpoint, 344 | use_scale_shift_norm, 345 | resblock_updown, 346 | use_fp16, 347 | ): 348 | model = sr_create_model( 349 | large_size, 350 | small_size, 351 | num_channels, 352 | num_res_blocks, 353 | learn_sigma=learn_sigma, 354 | class_cond=class_cond, 355 | use_checkpoint=use_checkpoint, 356 | attention_resolutions=attention_resolutions, 357 | num_heads=num_heads, 358 | num_head_channels=num_head_channels, 359 | num_heads_upsample=num_heads_upsample, 360 | use_scale_shift_norm=use_scale_shift_norm, 361 | dropout=dropout, 362 | resblock_updown=resblock_updown, 363 | use_fp16=use_fp16, 364 | ) 365 | diffusion = create_gaussian_diffusion( 366 | steps=diffusion_steps, 367 | learn_sigma=learn_sigma, 368 | noise_schedule=noise_schedule, 369 | use_kl=use_kl, 370 | predict_xstart=predict_xstart, 371 | rescale_timesteps=rescale_timesteps, 372 | rescale_learned_sigmas=rescale_learned_sigmas, 373 | timestep_respacing=timestep_respacing, 374 | ) 375 | return model, diffusion 376 | 377 | 378 | def sr_create_model( 379 | large_size, 380 | small_size, 381 | num_channels, 382 | num_res_blocks, 383 | learn_sigma, 384 | class_cond, 385 | use_checkpoint, 386 | attention_resolutions, 387 | num_heads, 388 | num_head_channels, 389 | num_heads_upsample, 390 | use_scale_shift_norm, 391 | dropout, 392 | resblock_updown, 393 | use_fp16, 394 | ): 395 | _ = small_size # hack to prevent unused variable 396 | 397 | if large_size == 512: 398 | channel_mult = (1, 1, 2, 2, 4, 4) 399 | elif large_size == 256: 400 | channel_mult = (1, 1, 2, 2, 4, 4) 401 | elif large_size == 64: 402 | channel_mult = (1, 2, 3, 4) 403 | elif large_size == 32: 404 | channel_mult = (1, 2, 4, 4) 405 | else: 406 | raise ValueError(f"unsupported large size: {large_size}") 407 | 408 | attention_ds = [] 409 | for res in attention_resolutions.split(","): 410 | attention_ds.append(large_size // int(res)) 411 | 412 | return SuperResModel( 413 | image_size=large_size, 414 | in_channels=3, 415 | model_channels=num_channels, 416 | out_channels=(3 if not learn_sigma else 6), 417 | num_res_blocks=num_res_blocks, 418 | attention_resolutions=tuple(attention_ds), 419 | dropout=dropout, 420 | channel_mult=channel_mult, 421 | num_classes=(NUM_CLASSES if class_cond else None), 422 | use_checkpoint=use_checkpoint, 423 | num_heads=num_heads, 424 | num_head_channels=num_head_channels, 425 | num_heads_upsample=num_heads_upsample, 426 | use_scale_shift_norm=use_scale_shift_norm, 427 | resblock_updown=resblock_updown, 428 | use_fp16=use_fp16, 429 | ) 430 | 431 | 432 | def create_gaussian_diffusion( 433 | *, 434 | steps=1000, 435 | learn_sigma=False, 436 | sigma_small=False, 437 | noise_schedule="linear", 438 | use_kl=False, 439 | predict_xstart=False, 440 | rescale_timesteps=False, 441 | rescale_learned_sigmas=False, 442 | timestep_respacing="", 443 | ): 444 | betas = gd.get_named_beta_schedule(noise_schedule, steps) 445 | if use_kl: 446 | loss_type = gd.LossType.RESCALED_KL 447 | elif rescale_learned_sigmas: 448 | loss_type = gd.LossType.RESCALED_MSE 449 | else: 450 | loss_type = gd.LossType.MSE 451 | if not timestep_respacing: 452 | timestep_respacing = [steps] 453 | return SpacedDiffusion( 454 | use_timesteps=space_timesteps(steps, timestep_respacing), 455 | betas=betas, 456 | model_mean_type=( 457 | gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X 458 | ), 459 | model_var_type=( 460 | ( 461 | gd.ModelVarType.FIXED_LARGE 462 | if not sigma_small 463 | else gd.ModelVarType.FIXED_SMALL 464 | ) 465 | if not learn_sigma 466 | else gd.ModelVarType.LEARNED_RANGE 467 | ), 468 | loss_type=loss_type, 469 | rescale_timesteps=rescale_timesteps, 470 | ) 471 | 472 | 473 | def add_dict_to_argparser(parser, default_dict): 474 | for k, v in default_dict.items(): 475 | v_type = type(v) 476 | if v is None: 477 | v_type = str 478 | elif isinstance(v, bool): 479 | v_type = str2bool 480 | parser.add_argument(f"--{k}", default=v, type=v_type) 481 | 482 | 483 | def args_to_dict(args, keys): 484 | return {k: getattr(args, k) for k in keys} 485 | 486 | 487 | def str2bool(v): 488 | """ 489 | https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse 490 | """ 491 | if isinstance(v, bool): 492 | return v 493 | if v.lower() in ("yes", "true", "t", "y", "1"): 494 | return True 495 | elif v.lower() in ("no", "false", "f", "n", "0"): 496 | return False 497 | else: 498 | raise argparse.ArgumentTypeError("boolean value expected") 499 | -------------------------------------------------------------------------------- /guided_diffusion/train_util.py: -------------------------------------------------------------------------------- 1 | import copy 2 | import functools 3 | import os 4 | 5 | import blobfile as bf 6 | import torch as th 7 | import torch.distributed as dist 8 | from torch.nn.parallel.distributed import DistributedDataParallel as DDP 9 | from torch.optim import AdamW 10 | 11 | from . import dist_util, logger 12 | from .fp16_util import MixedPrecisionTrainer 13 | from .nn import update_ema 14 | from .resample import LossAwareSampler, UniformSampler 15 | 16 | # For ImageNet experiments, this was a good default value. 17 | # We found that the lg_loss_scale quickly climbed to 18 | # 20-21 within the first ~1K steps of training. 19 | INITIAL_LOG_LOSS_SCALE = 20.0 20 | 21 | 22 | class TrainLoop: 23 | def __init__( 24 | self, 25 | *, 26 | model, 27 | diffusion, 28 | data, 29 | batch_size, 30 | microbatch, 31 | lr, 32 | ema_rate, 33 | log_interval, 34 | save_interval, 35 | resume_checkpoint, 36 | use_fp16=False, 37 | fp16_scale_growth=1e-3, 38 | schedule_sampler=None, 39 | weight_decay=0.0, 40 | lr_anneal_steps=0, 41 | lr_warmup_steps=0 42 | ): 43 | self.model = model 44 | self.diffusion = diffusion 45 | self.data = data 46 | self.batch_size = batch_size 47 | self.microbatch = microbatch if microbatch > 0 else batch_size 48 | self.lr = lr 49 | self.ema_rate = ( 50 | [ema_rate] 51 | if isinstance(ema_rate, float) 52 | else [float(x) for x in ema_rate.split(",")] 53 | ) 54 | self.log_interval = log_interval 55 | self.save_interval = save_interval 56 | self.resume_checkpoint = resume_checkpoint 57 | self.use_fp16 = use_fp16 58 | self.fp16_scale_growth = fp16_scale_growth 59 | self.schedule_sampler = schedule_sampler or UniformSampler(diffusion) 60 | self.weight_decay = weight_decay 61 | self.lr_anneal_steps = lr_anneal_steps 62 | self.lr_warmup_steps = lr_warmup_steps 63 | 64 | self.step = 0 65 | self.resume_step = 0 66 | self.global_batch = self.batch_size * dist.get_world_size() 67 | 68 | self.sync_cuda = th.cuda.is_available() 69 | 70 | self._load_and_sync_parameters() 71 | self.mp_trainer = MixedPrecisionTrainer( 72 | model=self.model, 73 | use_fp16=self.use_fp16, 74 | fp16_scale_growth=fp16_scale_growth, 75 | ) 76 | 77 | self.opt = AdamW( 78 | self.mp_trainer.master_params, lr=self.lr, weight_decay=self.weight_decay 79 | ) 80 | if self.resume_step: 81 | self._load_optimizer_state() 82 | # Model was resumed, either due to a restart or a checkpoint 83 | # being specified at the command line. 84 | self.ema_params = [ 85 | self._load_ema_parameters(rate) for rate in self.ema_rate 86 | ] 87 | else: 88 | self.ema_params = [ 89 | copy.deepcopy(self.mp_trainer.master_params) 90 | for _ in range(len(self.ema_rate)) 91 | ] 92 | 93 | if th.cuda.is_available(): 94 | self.use_ddp = True 95 | self.ddp_model = DDP( 96 | self.model, 97 | device_ids=[dist_util.dev()], 98 | output_device=dist_util.dev(), 99 | broadcast_buffers=False, 100 | bucket_cap_mb=128, 101 | find_unused_parameters=False, 102 | ) 103 | else: 104 | if dist.get_world_size() > 1: 105 | logger.warn( 106 | "Distributed training requires CUDA. " 107 | "Gradients will not be synchronized properly!" 108 | ) 109 | self.use_ddp = False 110 | self.ddp_model = self.model 111 | 112 | def _load_and_sync_parameters(self): 113 | resume_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 114 | 115 | if resume_checkpoint: 116 | self.resume_step = parse_resume_step_from_filename(resume_checkpoint) 117 | if dist.get_rank() == 0: 118 | if self.resume_step == 0: 119 | logger.log(f"loading model for the first time, will ignore missing layers if possible") 120 | state_dict = dist_util.load_state_dict( 121 | resume_checkpoint, map_location=dist_util.dev() 122 | ) 123 | 124 | model_state_dict = self.model.state_dict() 125 | for k in state_dict: 126 | if k in model_state_dict: 127 | if state_dict[k].shape != model_state_dict[k].shape: 128 | if k == 'input_blocks.0.0.weight': 129 | logger.info(f"Force load parameter: {k}, " 130 | f"required shape: {model_state_dict[k].shape}, " 131 | f"loaded shape: {state_dict[k].shape}") 132 | temp = th.rand(320,8,3,3) 133 | temp[:,:4,:,:] = state_dict[k] 134 | state_dict[k] = temp 135 | 136 | else: 137 | logger.info(f"Skip loading parameter: {k}, " 138 | f"required shape: {model_state_dict[k].shape}, " 139 | f"loaded shape: {state_dict[k].shape}") 140 | state_dict[k] = model_state_dict[k] 141 | if k.endswith('weight'): 142 | kb = k.replace('weight','bias') 143 | state_dict[kb] = model_state_dict[kb] 144 | else: 145 | logger.info(f"Dropping parameter {k}") 146 | 147 | self.model.load_state_dict(state_dict, strict=False) 148 | del state_dict 149 | del model_state_dict 150 | th.cuda.empty_cache() 151 | else: 152 | logger.log(f"loading model from checkpoint: {resume_checkpoint}...") 153 | self.model.load_state_dict( 154 | dist_util.load_state_dict( 155 | resume_checkpoint, map_location=dist_util.dev() 156 | ), strict=True 157 | ) 158 | 159 | dist_util.sync_params(self.model.parameters()) 160 | 161 | def _load_ema_parameters(self, rate): 162 | ema_params = copy.deepcopy(self.mp_trainer.master_params) 163 | 164 | main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 165 | ema_checkpoint = find_ema_checkpoint(main_checkpoint, self.resume_step, rate) 166 | #if ema_checkpoint: 167 | # if dist.get_rank() == 0: 168 | # logger.log(f"loading EMA from checkpoint: {ema_checkpoint}...") 169 | # state_dict = dist_util.load_state_dict( 170 | # ema_checkpoint, map_location=dist_util.dev() 171 | # ) 172 | # ema_params = self.mp_trainer.state_dict_to_master_params(state_dict) 173 | 174 | #dist_util.sync_params(ema_params) 175 | 176 | if ema_checkpoint: 177 | logger.log(f"loading EMA from checkpoint: {ema_checkpoint}...") 178 | state_dict = dist_util.load_state_dict( 179 | ema_checkpoint, map_location=dist_util.dev() 180 | ) 181 | ema_params = self.mp_trainer.state_dict_to_master_params(state_dict) 182 | 183 | return ema_params 184 | 185 | def _load_optimizer_state(self): 186 | main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 187 | opt_checkpoint = bf.join( 188 | bf.dirname(main_checkpoint), f"opt{self.resume_step:06}.pt" 189 | ) 190 | if bf.exists(opt_checkpoint): 191 | logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}") 192 | state_dict = dist_util.load_state_dict( 193 | opt_checkpoint, map_location=dist_util.dev() 194 | ) 195 | self.opt.load_state_dict(state_dict) 196 | 197 | def run_loop(self): 198 | while ( 199 | not self.lr_anneal_steps 200 | or self.step + self.resume_step < self.lr_anneal_steps 201 | ): 202 | batch, cond = next(self.data) 203 | self.run_step(batch, cond) 204 | if self.step % self.log_interval == 0: 205 | logger.dumpkvs() 206 | if self.step % self.save_interval == 0: 207 | self.save() 208 | # Run for a finite amount of time in integration tests. 209 | if os.environ.get("DIFFUSION_TRAINING_TEST", "") and self.step > 0: 210 | return 211 | self.step += 1 212 | # Save the last checkpoint if it wasn't already saved. 213 | if (self.step - 1) % self.save_interval != 0: 214 | self.save() 215 | 216 | def run_step(self, batch, cond): 217 | self.forward_backward(batch, cond) 218 | took_step = self.mp_trainer.optimize(self.opt) 219 | if took_step: 220 | self._update_ema() 221 | self._warmup_lr() 222 | self._anneal_lr() 223 | self.log_step() 224 | 225 | def forward_backward(self, batch, cond): 226 | self.mp_trainer.zero_grad() 227 | for i in range(0, batch.shape[0], self.microbatch): 228 | micro = batch[i : i + self.microbatch].to(dist_util.dev()) 229 | micro_cond = { 230 | k: v[i : i + self.microbatch].to(dist_util.dev()) 231 | for k, v in cond.items() 232 | } 233 | last_batch = (i + self.microbatch) >= batch.shape[0] 234 | t, weights = self.schedule_sampler.sample(micro.shape[0], dist_util.dev()) 235 | 236 | compute_losses = functools.partial( 237 | self.diffusion.training_losses, 238 | self.ddp_model, 239 | micro, 240 | t, 241 | model_kwargs=micro_cond, 242 | ) 243 | 244 | if last_batch or not self.use_ddp: 245 | losses = compute_losses() 246 | else: 247 | with self.ddp_model.no_sync(): 248 | losses = compute_losses() 249 | 250 | if isinstance(self.schedule_sampler, LossAwareSampler): 251 | self.schedule_sampler.update_with_local_losses( 252 | t, losses["loss"].detach() 253 | ) 254 | 255 | loss = (losses["loss"] * weights).mean() 256 | log_loss_dict( 257 | self.diffusion, t, {k: v * weights for k, v in losses.items()} 258 | ) 259 | self.mp_trainer.backward(loss) 260 | 261 | def _update_ema(self): 262 | for rate, params in zip(self.ema_rate, self.ema_params): 263 | update_ema(params, self.mp_trainer.master_params, rate=rate) 264 | 265 | def _anneal_lr(self): 266 | if not self.lr_anneal_steps: 267 | return 268 | frac_done = (self.step + self.resume_step) / self.lr_anneal_steps 269 | lr = self.lr * (1 - frac_done) 270 | for param_group in self.opt.param_groups: 271 | param_group["lr"] = lr 272 | 273 | def _warmup_lr(self): 274 | if not self.lr_warmup_steps: 275 | return 276 | frac_done = (self.step + self.resume_step) / self.lr_warmup_steps 277 | if frac_done > 1: 278 | return 279 | lr = self.lr * frac_done 280 | 281 | for param_group in self.opt.param_groups: 282 | param_group["lr"] = lr 283 | 284 | logger.log(f"setting lr to {lr}...") 285 | 286 | def log_step(self): 287 | logger.logkv("step", self.step + self.resume_step) 288 | logger.logkv("samples", (self.step + self.resume_step + 1) * self.global_batch) 289 | 290 | def save(self): 291 | def save_checkpoint(rate, params): 292 | state_dict = self.mp_trainer.master_params_to_state_dict(params) 293 | if dist.get_rank() == 0: 294 | logger.log(f"saving model {rate}...") 295 | if not rate: 296 | filename = f"model{(self.step+self.resume_step):06d}.pt" 297 | else: 298 | filename = f"ema_{rate}_{(self.step+self.resume_step):06d}.pt" 299 | with bf.BlobFile(bf.join(get_blob_logdir(), filename), "wb") as f: 300 | th.save(state_dict, f) 301 | 302 | save_checkpoint(0, self.mp_trainer.master_params) 303 | for rate, params in zip(self.ema_rate, self.ema_params): 304 | save_checkpoint(rate, params) 305 | 306 | if dist.get_rank() == 0: 307 | with bf.BlobFile( 308 | bf.join(get_blob_logdir(), f"opt{(self.step+self.resume_step):06d}.pt"), 309 | "wb", 310 | ) as f: 311 | th.save(self.opt.state_dict(), f) 312 | 313 | dist.barrier() 314 | 315 | 316 | def parse_resume_step_from_filename(filename): 317 | """ 318 | Parse filenames of the form path/to/modelNNNNNN.pt, where NNNNNN is the 319 | checkpoint's number of steps. 320 | """ 321 | split = filename.split("model") 322 | if len(split) < 2: 323 | return 0 324 | split1 = split[-1].split(".")[0] 325 | try: 326 | return int(split1) 327 | except ValueError: 328 | return 0 329 | 330 | 331 | def get_blob_logdir(): 332 | # You can change this to be a separate path to save checkpoints to 333 | # a blobstore or some external drive. 334 | return logger.get_dir() 335 | 336 | 337 | def find_resume_checkpoint(): 338 | # On your infrastructure, you may want to override this to automatically 339 | # discover the latest checkpoint on your blob storage, etc. 340 | return None 341 | 342 | 343 | def find_ema_checkpoint(main_checkpoint, step, rate): 344 | if main_checkpoint is None: 345 | return None 346 | filename = f"ema_{rate}_{(step):06d}.pt" 347 | path = bf.join(bf.dirname(main_checkpoint), filename) 348 | if bf.exists(path): 349 | return path 350 | return None 351 | 352 | 353 | def log_loss_dict(diffusion, ts, losses): 354 | for key, values in losses.items(): 355 | logger.logkv_mean(key, values.mean().item()) 356 | # Log the quantiles (four quartiles, in particular). 357 | for sub_t, sub_loss in zip(ts.cpu().numpy(), values.detach().cpu().numpy()): 358 | quartile = int(4 * sub_t / diffusion.num_timesteps) 359 | logger.logkv_mean(f"{key}_q{quartile}", sub_loss) 360 | -------------------------------------------------------------------------------- /kl.yaml: -------------------------------------------------------------------------------- 1 | model: 2 | base_learning_rate: 4.5e-6 3 | target: ldm.models.autoencoder.AutoencoderKL 4 | params: 5 | monitor: "val/rec_loss" 6 | embed_dim: 4 7 | lossconfig: 8 | target: torch.nn.Identity 9 | ddconfig: 10 | double_z: True 11 | z_channels: 4 12 | resolution: 256 13 | in_channels: 3 14 | out_ch: 3 15 | ch: 128 16 | ch_mult: [ 1,2,4,4 ] 17 | num_res_blocks: 2 18 | attn_resolutions: [ ] 19 | dropout: 0.0 20 | -------------------------------------------------------------------------------- /merge.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import sys 3 | 4 | model_path = sys.argv[1] 5 | diffusion_path = sys.argv[2] 6 | 7 | state = torch.load(model_path) 8 | diffusion = torch.load(diffusion_path) 9 | 10 | diffusion_prefix = 'model.diffusion_model.' 11 | 12 | for key in diffusion.keys(): 13 | state['state_dict'][diffusion_prefix + key] = diffusion[key] 14 | 15 | torch.save(state, 'model-merged.pt') 16 | -------------------------------------------------------------------------------- /model-card.md: -------------------------------------------------------------------------------- 1 | # Overview 2 | 3 | These are diffusion models and noised image classifiers described in the paper [Diffusion Models Beat GANs on Image Synthesis](https://arxiv.org/abs/2105.05233). 4 | Included in this release are the following models: 5 | 6 | * Noisy ImageNet classifiers at resolutions 64x64, 128x128, 256x256, 512x512 7 | * A class-unconditional ImageNet diffusion model at resolution 256x256 8 | * Class conditional ImageNet diffusion models at 64x64, 128x128, 256x256, 512x512 resolutions 9 | * Class-conditional ImageNet upsampling diffusion models: 64x64->256x256, 128x128->512x512 10 | * Diffusion models trained on three LSUN classes at 256x256 resolution: cat, horse, bedroom 11 | 12 | # Datasets 13 | 14 | All of the models we are releasing were either trained on the [ILSVRC 2012 subset of ImageNet](http://www.image-net.org/challenges/LSVRC/2012/) or on single classes of [LSUN](https://arxiv.org/abs/1506.03365). 15 | Here, we describe characteristics of these datasets which impact model behavior: 16 | 17 | **LSUN**: This dataset was collected in 2015 using a combination of human labeling (from Amazon Mechanical Turk) and automated data labeling. 18 | * Each of the three classes we consider contain over a million images. 19 | * The dataset creators found that the label accuracy was roughly 90% across the entire LSUN dataset when measured by trained experts. 20 | * Images are scraped from the internet, and LSUN cat images in particular tend to often follow a “meme” format. 21 | * We found that there are occasionally humans in these photos, including faces, especially within the cat class. 22 | 23 | **ILSVRC 2012 subset of ImageNet**: This dataset was curated in 2012 and consists of roughly one million images, each belonging to one of 1000 classes. 24 | * A large portion of the classes in this dataset are animals, plants, and other naturally-occurring objects. 25 | * Many images contain humans, although usually these humans aren’t reflected by the class label (e.g. the class “Tench, tinca tinca” contains many photos of people holding fish). 26 | 27 | # Performance 28 | 29 | These models are intended to generate samples consistent with their training distributions. 30 | This has been measured in terms of FID, Precision, and Recall. 31 | These metrics all rely on the representations of a [pre-trained Inception-V3 model](https://arxiv.org/abs/1512.00567), 32 | which was trained on ImageNet, and so is likely to focus more on the ImageNet classes (such as animals) than on other visual features (such as human faces). 33 | 34 | Qualitatively, the samples produced by these models often look highly realistic, especially when a diffusion model is combined with a noisy classifier. 35 | 36 | # Intended Use 37 | 38 | These models are intended to be used for research purposes only. 39 | In particular, they can be used as a baseline for generative modeling research, or as a starting point to build off of for such research. 40 | 41 | These models are not intended to be commercially deployed. 42 | Additionally, they are not intended to be used to create propaganda or offensive imagery. 43 | 44 | Before releasing these models, we probed their ability to ease the creation of targeted imagery, since doing so could be potentially harmful. 45 | We did this either by fine-tuning our ImageNet models on a target LSUN class, or through classifier guidance with publicly available [CLIP models](https://github.com/openai/CLIP). 46 | * To probe fine-tuning capabilities, we restricted our compute budget to roughly $100 and tried both standard fine-tuning, 47 | and a diffusion-specific approach where we train a specialized classifier for the LSUN class. The resulting FIDs were significantly worse than publicly available GAN models, indicating that fine-tuning an ImageNet diffusion model does not significantly lower the cost of image generation. 48 | * To probe guidance with CLIP, we tried two approaches for using pre-trained CLIP models for classifier guidance. Either we fed the noised image to CLIP directly and used its gradients, or we fed the diffusion model's denoised prediction to the CLIP model and differentiated through the whole process. In both cases, we found that it was difficult to recover information from the CLIP model, indicating that these diffusion models are unlikely to make it significantly easier to extract knowledge from CLIP compared to existing GAN models. 49 | 50 | # Limitations 51 | 52 | These models sometimes produce highly unrealistic outputs, particularly when generating images containing human faces. 53 | This may stem from ImageNet's emphasis on non-human objects. 54 | 55 | While classifier guidance can improve sample quality, it reduces diversity, resulting in some modes of the data distribution being underrepresented. 56 | This can potentially amplify existing biases in the training dataset such as gender and racial biases. 57 | 58 | Because ImageNet and LSUN contain images from the internet, they include photos of real people, and the model may have memorized some of the information contained in these photos. 59 | However, these images are already publicly available, and existing generative models trained on ImageNet have not demonstrated significant leakage of this information. -------------------------------------------------------------------------------- /scripts/classifier_sample.py: -------------------------------------------------------------------------------- 1 | """ 2 | Like image_sample.py, but use a noisy image classifier to guide the sampling 3 | process towards more realistic images. 4 | """ 5 | 6 | import argparse 7 | import os 8 | 9 | import numpy as np 10 | import torch as th 11 | import torch.distributed as dist 12 | import torch.nn.functional as F 13 | 14 | from guided_diffusion import dist_util, logger 15 | from guided_diffusion.script_util import ( 16 | NUM_CLASSES, 17 | model_and_diffusion_defaults, 18 | classifier_defaults, 19 | create_model_and_diffusion, 20 | create_classifier, 21 | add_dict_to_argparser, 22 | args_to_dict, 23 | ) 24 | 25 | 26 | def main(): 27 | args = create_argparser().parse_args() 28 | 29 | dist_util.setup_dist() 30 | logger.configure() 31 | 32 | logger.log("creating model and diffusion...") 33 | model, diffusion = create_model_and_diffusion( 34 | **args_to_dict(args, model_and_diffusion_defaults().keys()) 35 | ) 36 | model.load_state_dict( 37 | dist_util.load_state_dict(args.model_path, map_location="cpu") 38 | ) 39 | model.to(dist_util.dev()) 40 | if args.use_fp16: 41 | model.convert_to_fp16() 42 | model.eval() 43 | 44 | logger.log("loading classifier...") 45 | classifier = create_classifier(**args_to_dict(args, classifier_defaults().keys())) 46 | classifier.load_state_dict( 47 | dist_util.load_state_dict(args.classifier_path, map_location="cpu") 48 | ) 49 | classifier.to(dist_util.dev()) 50 | if args.classifier_use_fp16: 51 | classifier.convert_to_fp16() 52 | classifier.eval() 53 | 54 | def cond_fn(x, t, y=None): 55 | assert y is not None 56 | with th.enable_grad(): 57 | x_in = x.detach().requires_grad_(True) 58 | logits = classifier(x_in, t) 59 | log_probs = F.log_softmax(logits, dim=-1) 60 | selected = log_probs[range(len(logits)), y.view(-1)] 61 | return th.autograd.grad(selected.sum(), x_in)[0] * args.classifier_scale 62 | 63 | def model_fn(x, t, y=None): 64 | assert y is not None 65 | return model(x, t, y if args.class_cond else None) 66 | 67 | logger.log("sampling...") 68 | all_images = [] 69 | all_labels = [] 70 | while len(all_images) * args.batch_size < args.num_samples: 71 | model_kwargs = {} 72 | classes = th.randint( 73 | low=0, high=NUM_CLASSES, size=(args.batch_size,), device=dist_util.dev() 74 | ) 75 | model_kwargs["y"] = classes 76 | sample_fn = ( 77 | diffusion.p_sample_loop if not args.use_ddim else diffusion.ddim_sample_loop 78 | ) 79 | sample = sample_fn( 80 | model_fn, 81 | (args.batch_size, 3, args.image_size, args.image_size), 82 | clip_denoised=args.clip_denoised, 83 | model_kwargs=model_kwargs, 84 | cond_fn=cond_fn, 85 | device=dist_util.dev(), 86 | ) 87 | sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8) 88 | sample = sample.permute(0, 2, 3, 1) 89 | sample = sample.contiguous() 90 | 91 | gathered_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())] 92 | dist.all_gather(gathered_samples, sample) # gather not supported with NCCL 93 | all_images.extend([sample.cpu().numpy() for sample in gathered_samples]) 94 | gathered_labels = [th.zeros_like(classes) for _ in range(dist.get_world_size())] 95 | dist.all_gather(gathered_labels, classes) 96 | all_labels.extend([labels.cpu().numpy() for labels in gathered_labels]) 97 | logger.log(f"created {len(all_images) * args.batch_size} samples") 98 | 99 | arr = np.concatenate(all_images, axis=0) 100 | arr = arr[: args.num_samples] 101 | label_arr = np.concatenate(all_labels, axis=0) 102 | label_arr = label_arr[: args.num_samples] 103 | if dist.get_rank() == 0: 104 | shape_str = "x".join([str(x) for x in arr.shape]) 105 | out_path = os.path.join(logger.get_dir(), f"samples_{shape_str}.npz") 106 | logger.log(f"saving to {out_path}") 107 | np.savez(out_path, arr, label_arr) 108 | 109 | dist.barrier() 110 | logger.log("sampling complete") 111 | 112 | 113 | def create_argparser(): 114 | defaults = dict( 115 | clip_denoised=True, 116 | num_samples=10000, 117 | batch_size=16, 118 | use_ddim=False, 119 | model_path="", 120 | classifier_path="", 121 | classifier_scale=1.0, 122 | ) 123 | defaults.update(model_and_diffusion_defaults()) 124 | defaults.update(classifier_defaults()) 125 | parser = argparse.ArgumentParser() 126 | add_dict_to_argparser(parser, defaults) 127 | return parser 128 | 129 | 130 | if __name__ == "__main__": 131 | main() 132 | -------------------------------------------------------------------------------- /scripts/classifier_train.py: -------------------------------------------------------------------------------- 1 | """ 2 | Train a noised image classifier on ImageNet. 3 | """ 4 | 5 | import argparse 6 | import os 7 | 8 | import blobfile as bf 9 | import torch as th 10 | import torch.distributed as dist 11 | import torch.nn.functional as F 12 | from torch.nn.parallel.distributed import DistributedDataParallel as DDP 13 | from torch.optim import AdamW 14 | 15 | from guided_diffusion import dist_util, logger 16 | from guided_diffusion.fp16_util import MixedPrecisionTrainer 17 | from guided_diffusion.image_datasets import load_data 18 | from guided_diffusion.resample import create_named_schedule_sampler 19 | from guided_diffusion.script_util import ( 20 | add_dict_to_argparser, 21 | args_to_dict, 22 | classifier_and_diffusion_defaults, 23 | create_classifier_and_diffusion, 24 | ) 25 | from guided_diffusion.train_util import parse_resume_step_from_filename, log_loss_dict 26 | 27 | 28 | def main(): 29 | args = create_argparser().parse_args() 30 | 31 | dist_util.setup_dist() 32 | logger.configure() 33 | 34 | logger.log("creating model and diffusion...") 35 | model, diffusion = create_classifier_and_diffusion( 36 | **args_to_dict(args, classifier_and_diffusion_defaults().keys()) 37 | ) 38 | model.to(dist_util.dev()) 39 | if args.noised: 40 | schedule_sampler = create_named_schedule_sampler( 41 | args.schedule_sampler, diffusion 42 | ) 43 | 44 | resume_step = 0 45 | if args.resume_checkpoint: 46 | resume_step = parse_resume_step_from_filename(args.resume_checkpoint) 47 | if dist.get_rank() == 0: 48 | logger.log( 49 | f"loading model from checkpoint: {args.resume_checkpoint}... at {resume_step} step" 50 | ) 51 | model.load_state_dict( 52 | dist_util.load_state_dict( 53 | args.resume_checkpoint, map_location=dist_util.dev() 54 | ) 55 | ) 56 | 57 | # Needed for creating correct EMAs and fp16 parameters. 58 | dist_util.sync_params(model.parameters()) 59 | 60 | mp_trainer = MixedPrecisionTrainer( 61 | model=model, use_fp16=args.classifier_use_fp16, initial_lg_loss_scale=16.0 62 | ) 63 | 64 | model = DDP( 65 | model, 66 | device_ids=[dist_util.dev()], 67 | output_device=dist_util.dev(), 68 | broadcast_buffers=False, 69 | bucket_cap_mb=128, 70 | find_unused_parameters=False, 71 | ) 72 | 73 | logger.log("creating data loader...") 74 | data = load_data( 75 | data_dir=args.data_dir, 76 | batch_size=args.batch_size, 77 | image_size=args.image_size, 78 | class_cond=True, 79 | random_crop=True, 80 | ) 81 | if args.val_data_dir: 82 | val_data = load_data( 83 | data_dir=args.val_data_dir, 84 | batch_size=args.batch_size, 85 | image_size=args.image_size, 86 | class_cond=True, 87 | ) 88 | else: 89 | val_data = None 90 | 91 | logger.log(f"creating optimizer...") 92 | opt = AdamW(mp_trainer.master_params, lr=args.lr, weight_decay=args.weight_decay) 93 | if args.resume_checkpoint: 94 | opt_checkpoint = bf.join( 95 | bf.dirname(args.resume_checkpoint), f"opt{resume_step:06}.pt" 96 | ) 97 | logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}") 98 | opt.load_state_dict( 99 | dist_util.load_state_dict(opt_checkpoint, map_location=dist_util.dev()) 100 | ) 101 | 102 | logger.log("training classifier model...") 103 | 104 | def forward_backward_log(data_loader, prefix="train"): 105 | batch, extra = next(data_loader) 106 | labels = extra["y"].to(dist_util.dev()) 107 | 108 | batch = batch.to(dist_util.dev()) 109 | # Noisy images 110 | if args.noised: 111 | t, _ = schedule_sampler.sample(batch.shape[0], dist_util.dev()) 112 | batch = diffusion.q_sample(batch, t) 113 | else: 114 | t = th.zeros(batch.shape[0], dtype=th.long, device=dist_util.dev()) 115 | 116 | for i, (sub_batch, sub_labels, sub_t) in enumerate( 117 | split_microbatches(args.microbatch, batch, labels, t) 118 | ): 119 | logits = model(sub_batch, timesteps=sub_t) 120 | loss = F.cross_entropy(logits, sub_labels, reduction="none") 121 | 122 | losses = {} 123 | losses[f"{prefix}_loss"] = loss.detach() 124 | losses[f"{prefix}_acc@1"] = compute_top_k( 125 | logits, sub_labels, k=1, reduction="none" 126 | ) 127 | losses[f"{prefix}_acc@5"] = compute_top_k( 128 | logits, sub_labels, k=5, reduction="none" 129 | ) 130 | log_loss_dict(diffusion, sub_t, losses) 131 | del losses 132 | loss = loss.mean() 133 | if loss.requires_grad: 134 | if i == 0: 135 | mp_trainer.zero_grad() 136 | mp_trainer.backward(loss * len(sub_batch) / len(batch)) 137 | 138 | for step in range(args.iterations - resume_step): 139 | logger.logkv("step", step + resume_step) 140 | logger.logkv( 141 | "samples", 142 | (step + resume_step + 1) * args.batch_size * dist.get_world_size(), 143 | ) 144 | if args.anneal_lr: 145 | set_annealed_lr(opt, args.lr, (step + resume_step) / args.iterations) 146 | forward_backward_log(data) 147 | mp_trainer.optimize(opt) 148 | if val_data is not None and not step % args.eval_interval: 149 | with th.no_grad(): 150 | with model.no_sync(): 151 | model.eval() 152 | forward_backward_log(val_data, prefix="val") 153 | model.train() 154 | if not step % args.log_interval: 155 | logger.dumpkvs() 156 | if ( 157 | step 158 | and dist.get_rank() == 0 159 | and not (step + resume_step) % args.save_interval 160 | ): 161 | logger.log("saving model...") 162 | save_model(mp_trainer, opt, step + resume_step) 163 | 164 | if dist.get_rank() == 0: 165 | logger.log("saving model...") 166 | save_model(mp_trainer, opt, step + resume_step) 167 | dist.barrier() 168 | 169 | 170 | def set_annealed_lr(opt, base_lr, frac_done): 171 | lr = base_lr * (1 - frac_done) 172 | for param_group in opt.param_groups: 173 | param_group["lr"] = lr 174 | 175 | 176 | def save_model(mp_trainer, opt, step): 177 | if dist.get_rank() == 0: 178 | th.save( 179 | mp_trainer.master_params_to_state_dict(mp_trainer.master_params), 180 | os.path.join(logger.get_dir(), f"model{step:06d}.pt"), 181 | ) 182 | th.save(opt.state_dict(), os.path.join(logger.get_dir(), f"opt{step:06d}.pt")) 183 | 184 | 185 | def compute_top_k(logits, labels, k, reduction="mean"): 186 | _, top_ks = th.topk(logits, k, dim=-1) 187 | if reduction == "mean": 188 | return (top_ks == labels[:, None]).float().sum(dim=-1).mean().item() 189 | elif reduction == "none": 190 | return (top_ks == labels[:, None]).float().sum(dim=-1) 191 | 192 | 193 | def split_microbatches(microbatch, *args): 194 | bs = len(args[0]) 195 | if microbatch == -1 or microbatch >= bs: 196 | yield tuple(args) 197 | else: 198 | for i in range(0, bs, microbatch): 199 | yield tuple(x[i : i + microbatch] if x is not None else None for x in args) 200 | 201 | 202 | def create_argparser(): 203 | defaults = dict( 204 | data_dir="", 205 | val_data_dir="", 206 | noised=True, 207 | iterations=150000, 208 | lr=3e-4, 209 | weight_decay=0.0, 210 | anneal_lr=False, 211 | batch_size=4, 212 | microbatch=-1, 213 | schedule_sampler="uniform", 214 | resume_checkpoint="", 215 | log_interval=10, 216 | eval_interval=5, 217 | save_interval=10000, 218 | ) 219 | defaults.update(classifier_and_diffusion_defaults()) 220 | parser = argparse.ArgumentParser() 221 | add_dict_to_argparser(parser, defaults) 222 | return parser 223 | 224 | 225 | if __name__ == "__main__": 226 | main() 227 | -------------------------------------------------------------------------------- /scripts/classifier_train_stable.py: -------------------------------------------------------------------------------- 1 | """ 2 | Train a noised image classifier on ImageNet. 3 | """ 4 | 5 | import argparse 6 | import os 7 | 8 | import blobfile as bf 9 | import torch as th 10 | import torch.distributed as dist 11 | import torch.nn.functional as F 12 | from torch.nn.parallel.distributed import DistributedDataParallel as DDP 13 | from torch.optim import AdamW 14 | 15 | from guided_diffusion import dist_util, logger 16 | from guided_diffusion.fp16_util import MixedPrecisionTrainer 17 | from guided_diffusion.image_datasets_classifier import load_data 18 | from guided_diffusion.resample import create_named_schedule_sampler 19 | from guided_diffusion.script_util import ( 20 | add_dict_to_argparser, 21 | args_to_dict, 22 | classifier_and_diffusion_defaults, 23 | create_classifier_and_diffusion, 24 | ) 25 | from guided_diffusion.train_util import parse_resume_step_from_filename, log_loss_dict 26 | 27 | from ldm.util import instantiate_from_config 28 | 29 | from omegaconf import OmegaConf 30 | 31 | def set_requires_grad(model, value): 32 | for param in model.parameters(): 33 | param.requires_grad = value 34 | 35 | def main(): 36 | args = create_argparser().parse_args() 37 | 38 | dist_util.setup_dist() 39 | logger.configure() 40 | 41 | logger.log("loading vae...") 42 | 43 | kl_config = OmegaConf.load('kl.yaml') 44 | kl_sd = th.load(args.kl_model, map_location="cpu") 45 | 46 | encoder = instantiate_from_config(kl_config.model) 47 | encoder.load_state_dict(kl_sd, strict=True) 48 | 49 | encoder.to(dist_util.dev()) 50 | encoder.eval() 51 | encoder.requires_grad_(False) 52 | set_requires_grad(encoder, False) 53 | 54 | logger.log("creating model and diffusion...") 55 | model, diffusion = create_classifier_and_diffusion( 56 | **args_to_dict(args, classifier_and_diffusion_defaults().keys()) 57 | ) 58 | model.to(dist_util.dev()) 59 | if args.noised: 60 | schedule_sampler = create_named_schedule_sampler( 61 | args.schedule_sampler, diffusion 62 | ) 63 | 64 | resume_step = 0 65 | if args.resume_checkpoint: 66 | resume_step = parse_resume_step_from_filename(args.resume_checkpoint) 67 | if dist.get_rank() == 0: 68 | logger.log( 69 | f"loading model from checkpoint: {args.resume_checkpoint}... at {resume_step} step" 70 | ) 71 | model.load_state_dict( 72 | dist_util.load_state_dict( 73 | args.resume_checkpoint, map_location=dist_util.dev() 74 | ) 75 | ) 76 | 77 | # Needed for creating correct EMAs and fp16 parameters. 78 | dist_util.sync_params(model.parameters()) 79 | 80 | mp_trainer = MixedPrecisionTrainer( 81 | model=model, use_fp16=args.classifier_use_fp16, initial_lg_loss_scale=16.0 82 | ) 83 | 84 | model = DDP( 85 | model, 86 | device_ids=[dist_util.dev()], 87 | output_device=dist_util.dev(), 88 | broadcast_buffers=False, 89 | bucket_cap_mb=128, 90 | find_unused_parameters=False, 91 | ) 92 | 93 | logger.log("creating data loader...") 94 | data = load_data( 95 | good_data_dir=args.good_data_dir, 96 | bad_data_dir=args.bad_data_dir, 97 | batch_size=args.batch_size, 98 | image_size=args.actual_image_size, 99 | val_split=False, 100 | #val_num=10000, # first 10000 images used for validation 101 | ) 102 | 103 | val_data = None # not very useful.. 104 | 105 | #val_data = load_data( 106 | # good_data_dir=args.good_data_dir, 107 | # bad_data_dir=args.bad_data_dir, 108 | # batch_size=args.batch_size, 109 | # image_size=args.actual_image_size, 110 | # val_split=True, 111 | # val_num=10000, 112 | #) 113 | 114 | logger.log(f"creating optimizer...") 115 | opt = AdamW(mp_trainer.master_params, lr=args.lr, weight_decay=args.weight_decay) 116 | if args.resume_checkpoint: 117 | opt_checkpoint = bf.join( 118 | bf.dirname(args.resume_checkpoint), f"opt{resume_step:06}.pt" 119 | ) 120 | logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}") 121 | opt.load_state_dict( 122 | dist_util.load_state_dict(opt_checkpoint, map_location=dist_util.dev()) 123 | ) 124 | 125 | logger.log("training classifier model...") 126 | 127 | def forward_backward_log(data_loader, prefix="train"): 128 | batch, extra = next(data_loader) 129 | labels = extra["y"].to(dist_util.dev()) 130 | 131 | batch = batch.to(dist_util.dev()) 132 | 133 | batch = encoder.encode(batch).sample() 134 | batch *= 0.18215 135 | 136 | if args.classifier_use_fp16: 137 | batch = batch.half() 138 | 139 | # Noisy images 140 | if args.noised: 141 | t, _ = schedule_sampler.sample(batch.shape[0], dist_util.dev()) 142 | batch = diffusion.q_sample(batch, t) 143 | else: 144 | t = th.zeros(batch.shape[0], dtype=th.long, device=dist_util.dev()) 145 | 146 | for i, (sub_batch, sub_labels, sub_t) in enumerate( 147 | split_microbatches(args.microbatch, batch, labels, t) 148 | ): 149 | logits = model(sub_batch, timesteps=sub_t) 150 | loss = F.cross_entropy(logits, sub_labels, reduction="none") 151 | 152 | losses = {} 153 | losses[f"{prefix}_loss"] = loss.detach() 154 | # binary classification 155 | #losses[f"{prefix}_acc@1"] = compute_top_k( 156 | # logits, sub_labels, k=1, reduction="none" 157 | #) 158 | #losses[f"{prefix}_acc@5"] = compute_top_k( 159 | # logits, sub_labels, k=5, reduction="none" 160 | #) 161 | log_loss_dict(diffusion, sub_t, losses) 162 | del losses 163 | loss = loss.mean() 164 | if loss.requires_grad: 165 | if i == 0: 166 | mp_trainer.zero_grad() 167 | mp_trainer.backward(loss * len(sub_batch) / len(batch)) 168 | 169 | for step in range(args.iterations - resume_step): 170 | logger.logkv("step", step + resume_step) 171 | logger.logkv( 172 | "samples", 173 | (step + resume_step + 1) * args.batch_size * dist.get_world_size(), 174 | ) 175 | if args.anneal_lr: 176 | set_annealed_lr(opt, args.lr, (step + resume_step) / args.iterations) 177 | if args.lr_warmup_steps > 0: 178 | lr = set_warmup_lr(opt, args.lr, (step + resume_step) / args.lr_warmup_steps) 179 | logger.logkv("lr", lr) 180 | 181 | forward_backward_log(data) 182 | mp_trainer.optimize(opt) 183 | if val_data is not None and not step % args.eval_interval: 184 | with th.no_grad(): 185 | with model.no_sync(): 186 | model.eval() 187 | forward_backward_log(val_data, prefix="val") 188 | model.train() 189 | if not step % args.log_interval: 190 | logger.dumpkvs() 191 | if ( 192 | step 193 | and dist.get_rank() == 0 194 | and not (step + resume_step) % args.save_interval 195 | ): 196 | logger.log("saving model...") 197 | save_model(mp_trainer, opt, step + resume_step) 198 | 199 | if dist.get_rank() == 0: 200 | logger.log("saving model...") 201 | save_model(mp_trainer, opt, step + resume_step) 202 | dist.barrier() 203 | 204 | 205 | def set_annealed_lr(opt, base_lr, frac_done): 206 | lr = base_lr * (1 - frac_done) 207 | for param_group in opt.param_groups: 208 | param_group["lr"] = lr 209 | 210 | def set_warmup_lr(opt, base_lr, frac): 211 | if frac > 1: 212 | return base_lr 213 | 214 | lr = base_lr * frac 215 | 216 | if lr < 1e-9: 217 | lr = 1e-9 218 | 219 | for param_group in opt.param_groups: 220 | param_group["lr"] = lr 221 | 222 | return lr 223 | def save_model(mp_trainer, opt, step): 224 | if dist.get_rank() == 0: 225 | th.save( 226 | mp_trainer.master_params_to_state_dict(mp_trainer.master_params), 227 | os.path.join(logger.get_dir(), f"model{step:06d}.pt"), 228 | ) 229 | th.save(opt.state_dict(), os.path.join(logger.get_dir(), f"opt{step:06d}.pt")) 230 | 231 | 232 | def compute_top_k(logits, labels, k, reduction="mean"): 233 | _, top_ks = th.topk(logits, k, dim=-1) 234 | if reduction == "mean": 235 | return (top_ks == labels[:, None]).float().sum(dim=-1).mean().item() 236 | elif reduction == "none": 237 | return (top_ks == labels[:, None]).float().sum(dim=-1) 238 | 239 | 240 | def split_microbatches(microbatch, *args): 241 | bs = len(args[0]) 242 | if microbatch == -1 or microbatch >= bs: 243 | yield tuple(args) 244 | else: 245 | for i in range(0, bs, microbatch): 246 | yield tuple(x[i : i + microbatch] if x is not None else None for x in args) 247 | 248 | 249 | def create_argparser(): 250 | defaults = dict( 251 | good_data_dir="", 252 | bad_data_dir="", 253 | noised=True, 254 | iterations=150000, 255 | lr=3e-4, 256 | weight_decay=0.05, 257 | anneal_lr=False, 258 | batch_size=4, 259 | microbatch=-1, 260 | schedule_sampler="uniform", 261 | resume_checkpoint="", 262 | log_interval=10, 263 | eval_interval=10, 264 | save_interval=10000, 265 | kl_model=None, 266 | actual_image_size=512, 267 | lr_warmup_steps=0, 268 | ) 269 | defaults.update(classifier_and_diffusion_defaults()) 270 | parser = argparse.ArgumentParser() 271 | add_dict_to_argparser(parser, defaults) 272 | return parser 273 | 274 | 275 | if __name__ == "__main__": 276 | main() 277 | -------------------------------------------------------------------------------- /scripts/image_nll.py: -------------------------------------------------------------------------------- 1 | """ 2 | Approximate the bits/dimension for an image model. 3 | """ 4 | 5 | import argparse 6 | import os 7 | 8 | import numpy as np 9 | import torch.distributed as dist 10 | 11 | from guided_diffusion import dist_util, logger 12 | from guided_diffusion.image_datasets import load_data 13 | from guided_diffusion.script_util import ( 14 | model_and_diffusion_defaults, 15 | create_model_and_diffusion, 16 | add_dict_to_argparser, 17 | args_to_dict, 18 | ) 19 | 20 | 21 | def main(): 22 | args = create_argparser().parse_args() 23 | 24 | dist_util.setup_dist() 25 | logger.configure() 26 | 27 | logger.log("creating model and diffusion...") 28 | model, diffusion = create_model_and_diffusion( 29 | **args_to_dict(args, model_and_diffusion_defaults().keys()) 30 | ) 31 | model.load_state_dict( 32 | dist_util.load_state_dict(args.model_path, map_location="cpu") 33 | ) 34 | model.to(dist_util.dev()) 35 | model.eval() 36 | 37 | logger.log("creating data loader...") 38 | data = load_data( 39 | data_dir=args.data_dir, 40 | batch_size=args.batch_size, 41 | image_size=args.image_size, 42 | class_cond=args.class_cond, 43 | deterministic=True, 44 | ) 45 | 46 | logger.log("evaluating...") 47 | run_bpd_evaluation(model, diffusion, data, args.num_samples, args.clip_denoised) 48 | 49 | 50 | def run_bpd_evaluation(model, diffusion, data, num_samples, clip_denoised): 51 | all_bpd = [] 52 | all_metrics = {"vb": [], "mse": [], "xstart_mse": []} 53 | num_complete = 0 54 | while num_complete < num_samples: 55 | batch, model_kwargs = next(data) 56 | batch = batch.to(dist_util.dev()) 57 | model_kwargs = {k: v.to(dist_util.dev()) for k, v in model_kwargs.items()} 58 | minibatch_metrics = diffusion.calc_bpd_loop( 59 | model, batch, clip_denoised=clip_denoised, model_kwargs=model_kwargs 60 | ) 61 | 62 | for key, term_list in all_metrics.items(): 63 | terms = minibatch_metrics[key].mean(dim=0) / dist.get_world_size() 64 | dist.all_reduce(terms) 65 | term_list.append(terms.detach().cpu().numpy()) 66 | 67 | total_bpd = minibatch_metrics["total_bpd"] 68 | total_bpd = total_bpd.mean() / dist.get_world_size() 69 | dist.all_reduce(total_bpd) 70 | all_bpd.append(total_bpd.item()) 71 | num_complete += dist.get_world_size() * batch.shape[0] 72 | 73 | logger.log(f"done {num_complete} samples: bpd={np.mean(all_bpd)}") 74 | 75 | if dist.get_rank() == 0: 76 | for name, terms in all_metrics.items(): 77 | out_path = os.path.join(logger.get_dir(), f"{name}_terms.npz") 78 | logger.log(f"saving {name} terms to {out_path}") 79 | np.savez(out_path, np.mean(np.stack(terms), axis=0)) 80 | 81 | dist.barrier() 82 | logger.log("evaluation complete") 83 | 84 | 85 | def create_argparser(): 86 | defaults = dict( 87 | data_dir="", clip_denoised=True, num_samples=1000, batch_size=1, model_path="" 88 | ) 89 | defaults.update(model_and_diffusion_defaults()) 90 | parser = argparse.ArgumentParser() 91 | add_dict_to_argparser(parser, defaults) 92 | return parser 93 | 94 | 95 | if __name__ == "__main__": 96 | main() 97 | -------------------------------------------------------------------------------- /scripts/image_sample.py: -------------------------------------------------------------------------------- 1 | """ 2 | Generate a large batch of image samples from a model and save them as a large 3 | numpy array. This can be used to produce samples for FID evaluation. 4 | """ 5 | 6 | import argparse 7 | import os 8 | 9 | import numpy as np 10 | import torch as th 11 | import torch.distributed as dist 12 | 13 | from guided_diffusion import dist_util, logger 14 | from guided_diffusion.script_util import ( 15 | NUM_CLASSES, 16 | model_and_diffusion_defaults, 17 | create_model_and_diffusion, 18 | add_dict_to_argparser, 19 | args_to_dict, 20 | ) 21 | 22 | 23 | def main(): 24 | args = create_argparser().parse_args() 25 | 26 | dist_util.setup_dist() 27 | logger.configure() 28 | 29 | logger.log("creating model and diffusion...") 30 | model, diffusion = create_model_and_diffusion( 31 | **args_to_dict(args, model_and_diffusion_defaults().keys()) 32 | ) 33 | model.load_state_dict( 34 | dist_util.load_state_dict(args.model_path, map_location="cpu") 35 | ) 36 | model.to(dist_util.dev()) 37 | if args.use_fp16: 38 | model.convert_to_fp16() 39 | model.eval() 40 | 41 | logger.log("sampling...") 42 | all_images = [] 43 | all_labels = [] 44 | while len(all_images) * args.batch_size < args.num_samples: 45 | model_kwargs = {} 46 | if args.class_cond: 47 | classes = th.randint( 48 | low=0, high=NUM_CLASSES, size=(args.batch_size,), device=dist_util.dev() 49 | ) 50 | model_kwargs["y"] = classes 51 | sample_fn = ( 52 | diffusion.p_sample_loop if not args.use_ddim else diffusion.ddim_sample_loop 53 | ) 54 | sample = sample_fn( 55 | model, 56 | (args.batch_size, 3, args.image_size, args.image_size), 57 | clip_denoised=args.clip_denoised, 58 | model_kwargs=model_kwargs, 59 | ) 60 | sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8) 61 | sample = sample.permute(0, 2, 3, 1) 62 | sample = sample.contiguous() 63 | 64 | gathered_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())] 65 | dist.all_gather(gathered_samples, sample) # gather not supported with NCCL 66 | all_images.extend([sample.cpu().numpy() for sample in gathered_samples]) 67 | if args.class_cond: 68 | gathered_labels = [ 69 | th.zeros_like(classes) for _ in range(dist.get_world_size()) 70 | ] 71 | dist.all_gather(gathered_labels, classes) 72 | all_labels.extend([labels.cpu().numpy() for labels in gathered_labels]) 73 | logger.log(f"created {len(all_images) * args.batch_size} samples") 74 | 75 | arr = np.concatenate(all_images, axis=0) 76 | arr = arr[: args.num_samples] 77 | if args.class_cond: 78 | label_arr = np.concatenate(all_labels, axis=0) 79 | label_arr = label_arr[: args.num_samples] 80 | if dist.get_rank() == 0: 81 | shape_str = "x".join([str(x) for x in arr.shape]) 82 | out_path = os.path.join(logger.get_dir(), f"samples_{shape_str}.npz") 83 | logger.log(f"saving to {out_path}") 84 | if args.class_cond: 85 | np.savez(out_path, arr, label_arr) 86 | else: 87 | np.savez(out_path, arr) 88 | 89 | dist.barrier() 90 | logger.log("sampling complete") 91 | 92 | 93 | def create_argparser(): 94 | defaults = dict( 95 | clip_denoised=True, 96 | num_samples=10000, 97 | batch_size=16, 98 | use_ddim=False, 99 | model_path="", 100 | ) 101 | defaults.update(model_and_diffusion_defaults()) 102 | parser = argparse.ArgumentParser() 103 | add_dict_to_argparser(parser, defaults) 104 | return parser 105 | 106 | 107 | if __name__ == "__main__": 108 | main() 109 | -------------------------------------------------------------------------------- /scripts/image_train.py: -------------------------------------------------------------------------------- 1 | """ 2 | Train a diffusion model on images. 3 | """ 4 | 5 | import argparse 6 | 7 | from guided_diffusion import dist_util, logger 8 | from guided_diffusion.image_datasets import load_data 9 | from guided_diffusion.resample import create_named_schedule_sampler 10 | from guided_diffusion.script_util import ( 11 | model_and_diffusion_defaults, 12 | create_model_and_diffusion, 13 | args_to_dict, 14 | add_dict_to_argparser, 15 | ) 16 | from guided_diffusion.train_util import TrainLoop 17 | 18 | 19 | def main(): 20 | args = create_argparser().parse_args() 21 | 22 | dist_util.setup_dist() 23 | logger.configure() 24 | 25 | logger.log("creating model and diffusion...") 26 | model, diffusion = create_model_and_diffusion( 27 | **args_to_dict(args, model_and_diffusion_defaults().keys()) 28 | ) 29 | model.to(dist_util.dev()) 30 | schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion) 31 | 32 | logger.log("creating data loader...") 33 | data = load_data( 34 | data_dir=args.data_dir, 35 | batch_size=args.batch_size, 36 | image_size=args.image_size, 37 | class_cond=args.class_cond, 38 | ) 39 | logger.log("training...") 40 | TrainLoop( 41 | model=model, 42 | diffusion=diffusion, 43 | data=data, 44 | batch_size=args.batch_size, 45 | microbatch=args.microbatch, 46 | lr=args.lr, 47 | ema_rate=args.ema_rate, 48 | log_interval=args.log_interval, 49 | save_interval=args.save_interval, 50 | resume_checkpoint=args.resume_checkpoint, 51 | use_fp16=args.use_fp16, 52 | fp16_scale_growth=args.fp16_scale_growth, 53 | schedule_sampler=schedule_sampler, 54 | weight_decay=args.weight_decay, 55 | lr_anneal_steps=args.lr_anneal_steps, 56 | ).run_loop() 57 | 58 | 59 | def create_argparser(): 60 | defaults = dict( 61 | data_dir="", 62 | schedule_sampler="uniform", 63 | lr=1e-4, 64 | weight_decay=0.0, 65 | lr_anneal_steps=0, 66 | batch_size=1, 67 | microbatch=-1, # -1 disables microbatches 68 | ema_rate="0.9999", # comma-separated list of EMA values 69 | log_interval=10, 70 | save_interval=10000, 71 | resume_checkpoint="", 72 | use_fp16=False, 73 | fp16_scale_growth=1e-3, 74 | ) 75 | defaults.update(model_and_diffusion_defaults()) 76 | parser = argparse.ArgumentParser() 77 | add_dict_to_argparser(parser, defaults) 78 | return parser 79 | 80 | 81 | if __name__ == "__main__": 82 | main() 83 | -------------------------------------------------------------------------------- /scripts/image_train_inpaint.py: -------------------------------------------------------------------------------- 1 | """ 2 | Train a diffusion model on images. 3 | """ 4 | 5 | import argparse 6 | 7 | from guided_diffusion import dist_util, logger 8 | from guided_diffusion.image_text_datasets import load_data 9 | from guided_diffusion.resample import create_named_schedule_sampler 10 | from guided_diffusion.script_util import ( 11 | model_and_diffusion_defaults, 12 | create_model_and_diffusion, 13 | args_to_dict, 14 | add_dict_to_argparser, 15 | ) 16 | from guided_diffusion.train_util import TrainLoop 17 | import torch 18 | import random 19 | 20 | from transformers import CLIPTokenizer, CLIPTextModel 21 | from ldm.util import instantiate_from_config 22 | 23 | from omegaconf import OmegaConf 24 | from torchvision import transforms 25 | 26 | def set_requires_grad(model, value): 27 | for param in model.parameters(): 28 | param.requires_grad = value 29 | 30 | def main(): 31 | args = create_argparser().parse_args() 32 | 33 | dist_util.setup_dist() 34 | logger.configure() 35 | 36 | logger.log("loading vae...") 37 | 38 | kl_config = OmegaConf.load('kl.yaml') 39 | kl_sd = torch.load(args.kl_model, map_location="cpu") 40 | 41 | encoder = instantiate_from_config(kl_config.model) 42 | encoder.load_state_dict(kl_sd, strict=True) 43 | 44 | encoder.to(dist_util.dev()) 45 | encoder.eval() 46 | encoder.requires_grad_(False) 47 | set_requires_grad(encoder, False) 48 | 49 | 50 | logger.log("loading text encoder...") 51 | 52 | clip_version = 'openai/clip-vit-large-patch14' 53 | clip_tokenizer = CLIPTokenizer.from_pretrained(clip_version) 54 | clip_transformer = CLIPTextModel.from_pretrained(clip_version) 55 | clip_transformer.eval().requires_grad_(False).to(dist_util.dev()) 56 | 57 | logger.log("creating model and diffusion...") 58 | model, diffusion = create_model_and_diffusion( 59 | **args_to_dict(args, model_and_diffusion_defaults().keys()) 60 | ) 61 | 62 | model.to(dist_util.dev()) 63 | 64 | logger.log('total base parameters', sum(x.numel() for x in model.parameters())) 65 | 66 | schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion) 67 | 68 | logger.log("creating data loader...") 69 | data = load_latent_data( 70 | encoder, 71 | clip_tokenizer, 72 | clip_transformer, 73 | data_dir=args.data_dir, 74 | batch_size=args.batch_size, 75 | image_size=args.image_size, 76 | use_fp16=args.use_fp16, 77 | ) 78 | logger.log("training...") 79 | TrainLoop( 80 | model=model, 81 | diffusion=diffusion, 82 | data=data, 83 | batch_size=args.batch_size, 84 | microbatch=args.microbatch, 85 | lr=args.lr, 86 | ema_rate=args.ema_rate, 87 | log_interval=args.log_interval, 88 | save_interval=args.save_interval, 89 | resume_checkpoint=args.resume_checkpoint, 90 | use_fp16=args.use_fp16, 91 | fp16_scale_growth=args.fp16_scale_growth, 92 | schedule_sampler=schedule_sampler, 93 | weight_decay=args.weight_decay, 94 | lr_anneal_steps=args.lr_anneal_steps, 95 | ).run_loop() 96 | 97 | def load_latent_data(encoder, clip_tokenizer, clip_transformer, data_dir, batch_size, image_size, use_fp16): 98 | data = load_data( 99 | data_dir=data_dir, 100 | batch_size=batch_size, 101 | image_size=256, 102 | class_cond=False, 103 | ) 104 | 105 | blur = transforms.GaussianBlur(kernel_size=35, sigma=(0.1, 5)) 106 | 107 | for batch, model_kwargs, text in data: 108 | 109 | text = list(text) 110 | for i in range(len(text)): 111 | if random.randint(0,100) < 20: 112 | text[i] = '' 113 | 114 | text_encoded = clip_tokenizer(text, truncation=True, max_length=77, return_length=True, return_overflowing_tokens=False, padding="max_length", return_tensors="pt") 115 | text_tokens = text_encoded["input_ids"].to(dist_util.dev()) 116 | 117 | text_emb = clip_transformer(input_ids=text_tokens).last_hidden_state 118 | 119 | if use_fp16: 120 | text_emb = text_emb.half() 121 | 122 | model_kwargs["context"] = text_emb 123 | 124 | batch = batch.to(dist_util.dev()) 125 | emb = encoder.encode(batch).sample() 126 | emb *= 0.18215 127 | 128 | if use_fp16: 129 | emb = emb.half() 130 | 131 | emb_cond = emb.detach().clone() 132 | 133 | for i in range(batch.shape[0]): 134 | if random.randint(0,100) < 20: 135 | emb_cond[i,:,:,:] = 0 # unconditional 136 | else: 137 | if random.randint(0,100) < 50: # random mask 138 | mask = torch.randn(1, emb.shape[2], emb.shape[3]).to(dist_util.dev()) 139 | mask = blur(mask) 140 | mask = (mask > 0) 141 | mask = mask.repeat(4, 1, 1) 142 | mask = mask.float() 143 | emb_cond[i] *= mask 144 | else: 145 | # mask out 4 random rectangles 146 | for j in range(random.randint(1,4)): 147 | max_area = emb.shape[2]*emb.shape[3]//2 148 | 149 | w = random.randint(1,emb.shape[3]) 150 | h = random.randint(1,emb.shape[2]) 151 | if w*h > max_area: 152 | if random.randint(0,100) < 50: 153 | w = max_area//h 154 | else: 155 | h = max_area//w 156 | if w == emb.shape[3]: 157 | offsetx = 0 158 | else: 159 | offsetx = random.randint(0, emb.shape[3]-w) 160 | if h == emb.shape[2]: 161 | offsety = 0 162 | else: 163 | offsety = random.randint(0, emb.shape[2]-h) 164 | emb_cond[i,:, offsety:offsety+h, offsetx:offsetx+w] = 0 165 | 166 | model_kwargs["image_embed"] = emb_cond 167 | 168 | yield emb, model_kwargs 169 | 170 | def create_argparser(): 171 | defaults = dict( 172 | data_dir="", 173 | schedule_sampler="uniform", 174 | lr=1e-4, 175 | weight_decay=0.0, 176 | lr_anneal_steps=0, 177 | batch_size=1, 178 | microbatch=-1, # -1 disables microbatches 179 | ema_rate="0.9999", # comma-separated list of EMA values 180 | log_interval=10, 181 | save_interval=10000, 182 | resume_checkpoint="", 183 | use_fp16=False, 184 | fp16_scale_growth=1e-3, 185 | kl_model=None, 186 | actual_image_size=512, 187 | lr_warmup_steps=0, 188 | ) 189 | defaults.update(model_and_diffusion_defaults()) 190 | 191 | defaults['image_condition'] = True 192 | 193 | parser = argparse.ArgumentParser() 194 | add_dict_to_argparser(parser, defaults) 195 | return parser 196 | 197 | 198 | if __name__ == "__main__": 199 | main() 200 | -------------------------------------------------------------------------------- /scripts/image_train_stable.py: -------------------------------------------------------------------------------- 1 | """ 2 | Train a diffusion model on images. 3 | """ 4 | 5 | import argparse 6 | 7 | from guided_diffusion import dist_util, logger 8 | from guided_diffusion.image_text_datasets import load_data 9 | from guided_diffusion.resample import create_named_schedule_sampler 10 | from guided_diffusion.script_util import ( 11 | model_and_diffusion_defaults, 12 | create_model_and_diffusion, 13 | args_to_dict, 14 | add_dict_to_argparser, 15 | ) 16 | from guided_diffusion.train_util import TrainLoop 17 | import torch 18 | import random 19 | 20 | from transformers import CLIPTokenizer, CLIPTextModel 21 | from ldm.util import instantiate_from_config 22 | 23 | from omegaconf import OmegaConf 24 | 25 | def set_requires_grad(model, value): 26 | for param in model.parameters(): 27 | param.requires_grad = value 28 | 29 | def main(): 30 | args = create_argparser().parse_args() 31 | 32 | dist_util.setup_dist() 33 | logger.configure() 34 | 35 | logger.log("loading vae...") 36 | 37 | kl_config = OmegaConf.load('kl.yaml') 38 | kl_sd = torch.load(args.kl_model, map_location="cpu") 39 | 40 | encoder = instantiate_from_config(kl_config.model) 41 | encoder.load_state_dict(kl_sd, strict=True) 42 | 43 | encoder.to(dist_util.dev()) 44 | encoder.eval() 45 | encoder.requires_grad_(False) 46 | set_requires_grad(encoder, False) 47 | 48 | 49 | logger.log("loading text encoder...") 50 | 51 | clip_version = 'openai/clip-vit-large-patch14' 52 | clip_tokenizer = CLIPTokenizer.from_pretrained(clip_version) 53 | clip_transformer = CLIPTextModel.from_pretrained(clip_version) 54 | clip_transformer.eval().requires_grad_(False).to(dist_util.dev()) 55 | 56 | 57 | logger.log("creating model and diffusion...") 58 | model, diffusion = create_model_and_diffusion( 59 | **args_to_dict(args, model_and_diffusion_defaults().keys()) 60 | ) 61 | 62 | model.to(dist_util.dev()) 63 | 64 | logger.log('total base parameters', sum(x.numel() for x in model.parameters())) 65 | 66 | schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion) 67 | 68 | logger.log("creating data loader...") 69 | data = load_latent_data( 70 | encoder, 71 | clip_tokenizer, 72 | clip_transformer, 73 | data_dir=args.data_dir, 74 | batch_size=args.batch_size, 75 | image_size=args.actual_image_size, 76 | use_fp16=args.use_fp16 77 | ) 78 | logger.log("training...") 79 | TrainLoop( 80 | model=model, 81 | diffusion=diffusion, 82 | data=data, 83 | batch_size=args.batch_size, 84 | microbatch=args.microbatch, 85 | lr=args.lr, 86 | ema_rate=args.ema_rate, 87 | log_interval=args.log_interval, 88 | save_interval=args.save_interval, 89 | resume_checkpoint=args.resume_checkpoint, 90 | use_fp16=args.use_fp16, 91 | fp16_scale_growth=args.fp16_scale_growth, 92 | schedule_sampler=schedule_sampler, 93 | weight_decay=args.weight_decay, 94 | lr_anneal_steps=args.lr_anneal_steps, 95 | lr_warmup_steps=args.lr_warmup_steps, 96 | ).run_loop() 97 | 98 | def load_latent_data(encoder, clip_tokenizer, clip_transformer, data_dir, batch_size, image_size, use_fp16): 99 | data = load_data( 100 | data_dir=data_dir, 101 | batch_size=batch_size, 102 | image_size=image_size, 103 | class_cond=False, 104 | ) 105 | for batch, model_kwargs, text in data: 106 | 107 | text = list(text) 108 | for i in range(len(text)): 109 | if random.randint(0,100) < 20: 110 | text[i] = '' 111 | 112 | text_encoded = clip_tokenizer(text, truncation=True, max_length=77, return_length=True, return_overflowing_tokens=False, padding="max_length", return_tensors="pt") 113 | text_tokens = text_encoded["input_ids"].to(dist_util.dev()) 114 | 115 | text_emb = clip_transformer(input_ids=text_tokens).last_hidden_state 116 | 117 | if use_fp16: 118 | text_emb = text_emb.half() 119 | 120 | model_kwargs["context"] = text_emb 121 | 122 | batch = batch.to(dist_util.dev()) 123 | emb = encoder.encode(batch).sample() 124 | emb *= 0.18215 125 | 126 | if use_fp16: 127 | emb = emb.half() 128 | 129 | yield emb, model_kwargs 130 | 131 | def create_argparser(): 132 | defaults = dict( 133 | data_dir="", 134 | schedule_sampler="uniform", 135 | lr=1e-4, 136 | weight_decay=0.0, 137 | lr_anneal_steps=0, 138 | batch_size=1, 139 | microbatch=-1, # -1 disables microbatches 140 | ema_rate="0.9999", # comma-separated list of EMA values 141 | log_interval=10, 142 | save_interval=10000, 143 | resume_checkpoint="", 144 | use_fp16=False, 145 | fp16_scale_growth=1e-3, 146 | kl_model=None, 147 | actual_image_size=512, 148 | lr_warmup_steps=0, 149 | ) 150 | defaults.update(model_and_diffusion_defaults()) 151 | parser = argparse.ArgumentParser() 152 | add_dict_to_argparser(parser, defaults) 153 | return parser 154 | 155 | 156 | if __name__ == "__main__": 157 | main() 158 | -------------------------------------------------------------------------------- /scripts/image_train_super.py: -------------------------------------------------------------------------------- 1 | """ 2 | Train a diffusion model on images. 3 | """ 4 | 5 | import argparse 6 | 7 | from guided_diffusion import dist_util, logger 8 | from guided_diffusion.image_text_super import load_data 9 | from guided_diffusion.resample import create_named_schedule_sampler 10 | from guided_diffusion.script_util import ( 11 | model_and_diffusion_defaults, 12 | create_model_and_diffusion, 13 | args_to_dict, 14 | add_dict_to_argparser, 15 | ) 16 | from guided_diffusion.train_util import TrainLoop 17 | import torch 18 | import random 19 | 20 | from ldm.util import instantiate_from_config 21 | 22 | from omegaconf import OmegaConf 23 | 24 | import torchvision.transforms as T 25 | 26 | def set_requires_grad(model, value): 27 | for param in model.parameters(): 28 | param.requires_grad = value 29 | 30 | def main(): 31 | args = create_argparser().parse_args() 32 | 33 | dist_util.setup_dist() 34 | logger.configure() 35 | 36 | logger.log("loading vae...") 37 | 38 | kl_config = OmegaConf.load('kl.yaml') 39 | kl_sd = torch.load(args.kl_model, map_location="cpu") 40 | 41 | encoder = instantiate_from_config(kl_config.model) 42 | encoder.load_state_dict(kl_sd, strict=True) 43 | 44 | encoder.to(dist_util.dev()) 45 | encoder.eval() 46 | encoder.requires_grad_(False) 47 | set_requires_grad(encoder, False) 48 | 49 | logger.log("creating model and diffusion...") 50 | model, diffusion = create_model_and_diffusion( 51 | **args_to_dict(args, model_and_diffusion_defaults().keys()) 52 | ) 53 | 54 | model.to(dist_util.dev()) 55 | 56 | logger.log('total base parameters', sum(x.numel() for x in model.parameters())) 57 | 58 | schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion) 59 | 60 | logger.log("creating data loader...") 61 | data = load_latent_data( 62 | encoder, 63 | data_dir=args.data_dir, 64 | batch_size=args.batch_size, 65 | image_size=args.actual_image_size, 66 | ) 67 | 68 | logger.log("training...") 69 | TrainLoop( 70 | model=model, 71 | diffusion=diffusion, 72 | data=data, 73 | batch_size=args.batch_size, 74 | microbatch=args.microbatch, 75 | lr=args.lr, 76 | ema_rate=args.ema_rate, 77 | log_interval=args.log_interval, 78 | save_interval=args.save_interval, 79 | resume_checkpoint=args.resume_checkpoint, 80 | use_fp16=args.use_fp16, 81 | fp16_scale_growth=args.fp16_scale_growth, 82 | schedule_sampler=schedule_sampler, 83 | weight_decay=args.weight_decay, 84 | lr_anneal_steps=args.lr_anneal_steps, 85 | lr_warmup_steps=args.lr_warmup_steps, 86 | ).run_loop() 87 | 88 | def load_latent_data(encoder, data_dir, batch_size, image_size): 89 | data = load_data( 90 | data_dir=data_dir, 91 | batch_size=batch_size, 92 | image_size=image_size, 93 | class_cond=False, 94 | ) 95 | for batch, batch_small in data: 96 | model_kwargs = {} 97 | 98 | batch_small = batch_small.to(dist_util.dev()) 99 | emb = encoder.encode(batch_small).sample().half() 100 | emb *= 0.18215 101 | 102 | model_kwargs["super_res_embed"] = emb 103 | 104 | yield batch, model_kwargs 105 | 106 | def create_argparser(): 107 | defaults = dict( 108 | data_dir="", 109 | schedule_sampler="uniform", 110 | lr=1e-4, 111 | weight_decay=0.0, 112 | lr_anneal_steps=0, 113 | batch_size=1, 114 | microbatch=-1, # -1 disables microbatches 115 | ema_rate="0.9999", # comma-separated list of EMA values 116 | log_interval=10, 117 | save_interval=10000, 118 | resume_checkpoint="", 119 | use_fp16=False, 120 | fp16_scale_growth=1e-3, 121 | kl_model=None, 122 | actual_image_size=1024, 123 | lr_warmup_steps=0, 124 | ) 125 | defaults.update(model_and_diffusion_defaults()) 126 | 127 | defaults['image_condition'] = False 128 | defaults['super_res_condition'] = True 129 | defaults['context_dim'] = None 130 | defaults['use_new_attention_order'] = False 131 | defaults['use_spatial_transformer'] = False 132 | 133 | parser = argparse.ArgumentParser() 134 | add_dict_to_argparser(parser, defaults) 135 | return parser 136 | 137 | 138 | if __name__ == "__main__": 139 | main() 140 | -------------------------------------------------------------------------------- /scripts/super_res_sample.py: -------------------------------------------------------------------------------- 1 | """ 2 | Generate a large batch of samples from a super resolution model, given a batch 3 | of samples from a regular model from image_sample.py. 4 | """ 5 | 6 | import argparse 7 | import os 8 | 9 | import blobfile as bf 10 | import numpy as np 11 | import torch as th 12 | import torch.distributed as dist 13 | 14 | from guided_diffusion import dist_util, logger 15 | from guided_diffusion.script_util import ( 16 | sr_model_and_diffusion_defaults, 17 | sr_create_model_and_diffusion, 18 | args_to_dict, 19 | add_dict_to_argparser, 20 | ) 21 | 22 | 23 | def main(): 24 | args = create_argparser().parse_args() 25 | 26 | dist_util.setup_dist() 27 | logger.configure() 28 | 29 | logger.log("creating model...") 30 | model, diffusion = sr_create_model_and_diffusion( 31 | **args_to_dict(args, sr_model_and_diffusion_defaults().keys()) 32 | ) 33 | model.load_state_dict( 34 | dist_util.load_state_dict(args.model_path, map_location="cpu") 35 | ) 36 | model.to(dist_util.dev()) 37 | if args.use_fp16: 38 | model.convert_to_fp16() 39 | model.eval() 40 | 41 | logger.log("loading data...") 42 | data = load_data_for_worker(args.base_samples, args.batch_size, args.class_cond) 43 | 44 | logger.log("creating samples...") 45 | all_images = [] 46 | while len(all_images) * args.batch_size < args.num_samples: 47 | model_kwargs = next(data) 48 | model_kwargs = {k: v.to(dist_util.dev()) for k, v in model_kwargs.items()} 49 | sample = diffusion.p_sample_loop( 50 | model, 51 | (args.batch_size, 3, args.large_size, args.large_size), 52 | clip_denoised=args.clip_denoised, 53 | model_kwargs=model_kwargs, 54 | ) 55 | sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8) 56 | sample = sample.permute(0, 2, 3, 1) 57 | sample = sample.contiguous() 58 | 59 | all_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())] 60 | dist.all_gather(all_samples, sample) # gather not supported with NCCL 61 | for sample in all_samples: 62 | all_images.append(sample.cpu().numpy()) 63 | logger.log(f"created {len(all_images) * args.batch_size} samples") 64 | 65 | arr = np.concatenate(all_images, axis=0) 66 | arr = arr[: args.num_samples] 67 | if dist.get_rank() == 0: 68 | shape_str = "x".join([str(x) for x in arr.shape]) 69 | out_path = os.path.join(logger.get_dir(), f"samples_{shape_str}.npz") 70 | logger.log(f"saving to {out_path}") 71 | np.savez(out_path, arr) 72 | 73 | dist.barrier() 74 | logger.log("sampling complete") 75 | 76 | 77 | def load_data_for_worker(base_samples, batch_size, class_cond): 78 | with bf.BlobFile(base_samples, "rb") as f: 79 | obj = np.load(f) 80 | image_arr = obj["arr_0"] 81 | if class_cond: 82 | label_arr = obj["arr_1"] 83 | rank = dist.get_rank() 84 | num_ranks = dist.get_world_size() 85 | buffer = [] 86 | label_buffer = [] 87 | while True: 88 | for i in range(rank, len(image_arr), num_ranks): 89 | buffer.append(image_arr[i]) 90 | if class_cond: 91 | label_buffer.append(label_arr[i]) 92 | if len(buffer) == batch_size: 93 | batch = th.from_numpy(np.stack(buffer)).float() 94 | batch = batch / 127.5 - 1.0 95 | batch = batch.permute(0, 3, 1, 2) 96 | res = dict(low_res=batch) 97 | if class_cond: 98 | res["y"] = th.from_numpy(np.stack(label_buffer)) 99 | yield res 100 | buffer, label_buffer = [], [] 101 | 102 | 103 | def create_argparser(): 104 | defaults = dict( 105 | clip_denoised=True, 106 | num_samples=10000, 107 | batch_size=16, 108 | use_ddim=False, 109 | base_samples="", 110 | model_path="", 111 | ) 112 | defaults.update(sr_model_and_diffusion_defaults()) 113 | parser = argparse.ArgumentParser() 114 | add_dict_to_argparser(parser, defaults) 115 | return parser 116 | 117 | 118 | if __name__ == "__main__": 119 | main() 120 | -------------------------------------------------------------------------------- /scripts/super_res_train.py: -------------------------------------------------------------------------------- 1 | """ 2 | Train a super-resolution model. 3 | """ 4 | 5 | import argparse 6 | 7 | import torch.nn.functional as F 8 | 9 | from guided_diffusion import dist_util, logger 10 | from guided_diffusion.image_datasets import load_data 11 | from guided_diffusion.resample import create_named_schedule_sampler 12 | from guided_diffusion.script_util import ( 13 | sr_model_and_diffusion_defaults, 14 | sr_create_model_and_diffusion, 15 | args_to_dict, 16 | add_dict_to_argparser, 17 | ) 18 | from guided_diffusion.train_util import TrainLoop 19 | 20 | 21 | def main(): 22 | args = create_argparser().parse_args() 23 | 24 | dist_util.setup_dist() 25 | logger.configure() 26 | 27 | logger.log("creating model...") 28 | model, diffusion = sr_create_model_and_diffusion( 29 | **args_to_dict(args, sr_model_and_diffusion_defaults().keys()) 30 | ) 31 | model.to(dist_util.dev()) 32 | schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion) 33 | 34 | logger.log("creating data loader...") 35 | data = load_superres_data( 36 | args.data_dir, 37 | args.batch_size, 38 | large_size=args.large_size, 39 | small_size=args.small_size, 40 | class_cond=args.class_cond, 41 | ) 42 | 43 | logger.log("training...") 44 | TrainLoop( 45 | model=model, 46 | diffusion=diffusion, 47 | data=data, 48 | batch_size=args.batch_size, 49 | microbatch=args.microbatch, 50 | lr=args.lr, 51 | ema_rate=args.ema_rate, 52 | log_interval=args.log_interval, 53 | save_interval=args.save_interval, 54 | resume_checkpoint=args.resume_checkpoint, 55 | use_fp16=args.use_fp16, 56 | fp16_scale_growth=args.fp16_scale_growth, 57 | schedule_sampler=schedule_sampler, 58 | weight_decay=args.weight_decay, 59 | lr_anneal_steps=args.lr_anneal_steps, 60 | ).run_loop() 61 | 62 | 63 | def load_superres_data(data_dir, batch_size, large_size, small_size, class_cond=False): 64 | data = load_data( 65 | data_dir=data_dir, 66 | batch_size=batch_size, 67 | image_size=large_size, 68 | class_cond=class_cond, 69 | ) 70 | for large_batch, model_kwargs in data: 71 | model_kwargs["low_res"] = F.interpolate(large_batch, small_size, mode="area") 72 | yield large_batch, model_kwargs 73 | 74 | 75 | def create_argparser(): 76 | defaults = dict( 77 | data_dir="", 78 | schedule_sampler="uniform", 79 | lr=1e-4, 80 | weight_decay=0.0, 81 | lr_anneal_steps=0, 82 | batch_size=1, 83 | microbatch=-1, 84 | ema_rate="0.9999", 85 | log_interval=10, 86 | save_interval=10000, 87 | resume_checkpoint="", 88 | use_fp16=False, 89 | fp16_scale_growth=1e-3, 90 | ) 91 | defaults.update(sr_model_and_diffusion_defaults()) 92 | parser = argparse.ArgumentParser() 93 | add_dict_to_argparser(parser, defaults) 94 | return parser 95 | 96 | 97 | if __name__ == "__main__": 98 | main() 99 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | setup( 4 | name="guided-diffusion", 5 | py_modules=["guided_diffusion"], 6 | install_requires=["blobfile>=1.0.5", "torch", "tqdm"], 7 | ) 8 | -------------------------------------------------------------------------------- /split.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import sys 3 | 4 | model_path = sys.argv[1] 5 | 6 | state = torch.load(model_path)['state_dict'] 7 | 8 | model = {} 9 | vae = {} 10 | 11 | model_prefix = 'model.diffusion_model.' 12 | vae_prefix = 'first_stage_model.' 13 | 14 | for key in state.keys(): 15 | if key.startswith(model_prefix): 16 | model[key[len(model_prefix):]] = state[key] 17 | elif key.startswith(vae_prefix): 18 | vae[key[len(vae_prefix):]] = state[key] 19 | 20 | torch.save(model, 'diffusion.pt') 21 | torch.save(vae, 'kl.pt') 22 | -------------------------------------------------------------------------------- /super_large.py: -------------------------------------------------------------------------------- 1 | import gc 2 | import io 3 | import math 4 | import sys 5 | 6 | from PIL import Image, ImageOps 7 | import requests 8 | import torch 9 | from torch import nn 10 | from torch.nn import functional as F 11 | from torchvision import transforms 12 | from torchvision.transforms import functional as TF 13 | from tqdm.notebook import tqdm 14 | 15 | import numpy as np 16 | 17 | from guided_diffusion.script_util import create_model_and_diffusion, model_and_diffusion_defaults 18 | 19 | from omegaconf import OmegaConf 20 | from ldm.util import instantiate_from_config 21 | 22 | from einops import rearrange 23 | from math import log2, sqrt 24 | 25 | import argparse 26 | import pickle 27 | 28 | import os 29 | 30 | # argument parsing 31 | 32 | parser = argparse.ArgumentParser() 33 | 34 | parser.add_argument('--model_path', type=str, default = 'upscale.pt', 35 | help='path to the diffusion model') 36 | 37 | parser.add_argument('--kl_path', type=str, default = 'kl.pt', 38 | help='path to the LDM first stage model') 39 | 40 | parser.add_argument('--image', type = str, required = True, default = '', 41 | help='path to image (npy containing latent embeddings or image file)') 42 | 43 | parser.add_argument('--text', type = str, required = False, default = '', 44 | help='your text prompt (not necessary). If provided, use clip guidance') 45 | 46 | parser.add_argument('--cutn', type = int, default = 32, required = False, 47 | help='Number of cuts for clip guidance') 48 | 49 | parser.add_argument('--clip_guidance_scale', type = int, default = 5000, required = False, 50 | help='Amount of clip guidance') 51 | 52 | parser.add_argument('--tv_scale', type = int, default = 0, required = False, 53 | help='Controls the smoothness of the final output') 54 | 55 | parser.add_argument('--range_scale', type = int, default = 0, required = False, 56 | help='Controls how far out of range RGB values are allowed to be') 57 | 58 | parser.add_argument('--init_image', type=str, required = False, default = None, 59 | help='init image to use') 60 | 61 | parser.add_argument('--skip_timesteps', type=int, required = False, default = 0, 62 | help='how many diffusion steps are gonna be skipped') 63 | 64 | parser.add_argument('--prefix', type = str, required = False, default = 'super_', 65 | help='prefix for output files') 66 | 67 | parser.add_argument('--num_batches', type = int, default = 1, required = False, 68 | help='number of batches') 69 | 70 | parser.add_argument('--batch_size', type = int, default = 1, required = False, 71 | help='batch size') 72 | 73 | parser.add_argument('--width', type = int, default = 0, required = False, 74 | help='image size of output (multiple of 8)') 75 | 76 | parser.add_argument('--height', type = int, default = 0, required = False, 77 | help='image size of output (multiple of 8)') 78 | 79 | parser.add_argument('--seed', type = int, default=-1, required = False, 80 | help='random seed') 81 | 82 | parser.add_argument('--steps', type = int, default = 0, required = False, 83 | help='number of diffusion steps') 84 | 85 | parser.add_argument('--cpu', dest='cpu', action='store_true') 86 | 87 | parser.add_argument('--guide', dest='guide', action='store_true') 88 | 89 | parser.add_argument('--avg', dest='avg', action='store_true') 90 | 91 | parser.add_argument('--clip_score', dest='clip_score', action='store_true') 92 | 93 | parser.add_argument('--ddim', dest='ddim', action='store_true') # turn on to use 50 step ddim 94 | 95 | parser.add_argument('--ddpm', dest='ddpm', action='store_true') # turn on to use 50 step ddim 96 | 97 | args = parser.parse_args() 98 | 99 | def fetch(url_or_path): 100 | if str(url_or_path).startswith('http://') or str(url_or_path).startswith('https://'): 101 | r = requests.get(url_or_path) 102 | r.raise_for_status() 103 | fd = io.BytesIO() 104 | fd.write(r.content) 105 | fd.seek(0) 106 | return fd 107 | return open(url_or_path, 'rb') 108 | 109 | def parse_prompt(prompt): 110 | if prompt.startswith('http://') or prompt.startswith('https://'): 111 | vals = prompt.rsplit(':', 2) 112 | vals = [vals[0] + ':' + vals[1], *vals[2:]] 113 | else: 114 | vals = prompt.rsplit(':', 1) 115 | vals = vals + ['', '1'][len(vals):] 116 | return vals[0], float(vals[1]) 117 | 118 | class MakeCutouts(nn.Module): 119 | def __init__(self, cut_size, cutn, cut_pow=1.): 120 | super().__init__() 121 | 122 | self.cut_size = cut_size 123 | self.cutn = cutn 124 | self.cut_pow = cut_pow 125 | 126 | def forward(self, input): 127 | sideY, sideX = input.shape[2:4] 128 | max_size = min(sideX, sideY) 129 | min_size = min(sideX, sideY, self.cut_size) 130 | cutouts = [] 131 | for _ in range(self.cutn): 132 | size = int(torch.rand([])**self.cut_pow * (max_size - min_size) + min_size) 133 | offsetx = torch.randint(0, sideX - size + 1, ()) 134 | offsety = torch.randint(0, sideY - size + 1, ()) 135 | cutout = input[:, :, offsety:offsety + size, offsetx:offsetx + size] 136 | cutouts.append(F.adaptive_avg_pool2d(cutout, self.cut_size)) 137 | return torch.cat(cutouts) 138 | 139 | 140 | def spherical_dist_loss(x, y): 141 | x = F.normalize(x, dim=-1) 142 | y = F.normalize(y, dim=-1) 143 | return (x - y).norm(dim=-1).div(2).arcsin().pow(2).mul(2) 144 | 145 | 146 | def tv_loss(input): 147 | """L2 total variation loss, as in Mahendran et al.""" 148 | input = F.pad(input, (0, 1, 0, 1), 'replicate') 149 | x_diff = input[..., :-1, 1:] - input[..., :-1, :-1] 150 | y_diff = input[..., 1:, :-1] - input[..., :-1, :-1] 151 | return (x_diff**2 + y_diff**2).mean([1, 2, 3]) 152 | 153 | 154 | def range_loss(input): 155 | return (input - input.clamp(-1, 1)).pow(2).mean([1, 2, 3]) 156 | 157 | device = torch.device('cuda:0' if (torch.cuda.is_available() and not args.cpu) else 'cpu') 158 | print('Using device:', device) 159 | 160 | model_state_dict = torch.load(args.model_path, map_location='cpu') 161 | 162 | model_params = { 163 | 'attention_resolutions': '64,32', 164 | 'class_cond': False, 165 | 'diffusion_steps': 1000, 166 | 'rescale_timesteps': True, 167 | 'timestep_respacing': '15', 168 | 'image_size': 1024, 169 | 'learn_sigma': True, 170 | 'noise_schedule': 'linear_openai', 171 | 'num_channels': 256, 172 | 'num_head_channels': 64, 173 | 'num_res_blocks': 2, 174 | 'resblock_updown': True, 175 | 'use_fp16': True, 176 | 'use_scale_shift_norm': True, 177 | 'clip_embed_dim': None, 178 | 'image_condition': False, 179 | 'super_res_condition': True, 180 | 'context_dim': None, 181 | 'use_spatial_transformer': False, 182 | } 183 | 184 | if args.ddpm: 185 | model_params['timestep_respacing'] = '1000' 186 | if args.ddim: 187 | if args.steps: 188 | model_params['timestep_respacing'] = 'ddim'+str(args.steps) 189 | else: 190 | model_params['timestep_respacing'] = 'ddim50' 191 | elif args.steps: 192 | model_params['timestep_respacing'] = str(args.steps) 193 | 194 | model_config = model_and_diffusion_defaults() 195 | model_config.update(model_params) 196 | 197 | if args.cpu: 198 | model_config['use_fp16'] = False 199 | 200 | # Load models 201 | model, diffusion = create_model_and_diffusion(**model_config) 202 | model.load_state_dict(model_state_dict, strict=True) 203 | model.requires_grad_(True if args.text or args.guide else False).eval().to(device) 204 | 205 | if model_config['use_fp16']: 206 | model.convert_to_fp16() 207 | else: 208 | model.convert_to_fp32() 209 | 210 | def set_requires_grad(model, value): 211 | for param in model.parameters(): 212 | param.requires_grad = value 213 | 214 | kl_config = OmegaConf.load('kl.yaml') 215 | kl_sd = torch.load(args.kl_path, map_location="cpu") 216 | 217 | ldm = instantiate_from_config(kl_config.model) 218 | ldm.load_state_dict(kl_sd, strict=True) 219 | 220 | ldm.to(device) 221 | ldm.eval() 222 | ldm.requires_grad_(False) 223 | set_requires_grad(ldm, False) 224 | 225 | if args.image.endswith('.npy'): 226 | image = np.load(args.image) 227 | im_tensor = torch.from_numpy(image).unsqueeze(0).to(device) 228 | input_image = ldm.decode(im_tensor).clamp(-1,1) 229 | input_image_small = F.interpolate(input_image, size=(im_tensor.shape[3]*8, im_tensor.shape[2]*8), mode='bilinear') 230 | args.width = im_tensor.shape[3]*16 231 | args.height = im_tensor.shape[2]*16 232 | else: 233 | image = Image.open(args.image).convert('RGB') 234 | 235 | if args.width == 0 and args.height == 0: 236 | crop_width = math.floor(image.width/64)*64 237 | crop_height = math.floor(image.height/64)*64 238 | 239 | left = (image.width-crop_width)//2 240 | top = (image.height-crop_height)//2 241 | 242 | image = image.crop((left, top, left+crop_width, top+crop_height)) 243 | 244 | args.width = 2*crop_width 245 | args.height = 2*crop_height 246 | else: 247 | image = ImageOps.fit(image, (args.width//2, args.height//2)) 248 | 249 | input_image = TF.to_tensor(image).to(device).unsqueeze(0) 250 | input_image = input_image * 2 - 1 251 | 252 | im_tensor = ldm.encode(input_image).sample().to(device) 253 | input_image_small = F.interpolate(input_image, size=(im_tensor.shape[3]*4, im_tensor.shape[2]*4), mode='bilinear') 254 | 255 | if args.text: 256 | import clip 257 | 258 | # if text is given, use clip guidance 259 | for name, param in model.named_parameters(): 260 | if 'qkv' in name or 'norm' in name or 'proj' in name: 261 | param.requires_grad_() 262 | 263 | clip_model, clip_preprocess = clip.load('ViT-L/14', jit=False) 264 | clip_model.eval().requires_grad_(False).to(device) 265 | clip_size = clip_model.visual.input_resolution 266 | 267 | normalize = transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711]) 268 | make_cutouts = MakeCutouts(clip_size, args.cutn) 269 | 270 | side_x = args.width 271 | side_y = args.height 272 | 273 | target_embeds, weights = [], [] 274 | 275 | prompts = args.text.split('|') 276 | 277 | for prompt in prompts: 278 | txt, weight = parse_prompt(prompt) 279 | target_embeds.append(clip_model.encode_text(clip.tokenize(prompt).to(device)).float()) 280 | weights.append(weight) 281 | 282 | target_embeds = torch.cat(target_embeds) 283 | weights = torch.tensor(weights, device=device) 284 | if weights.sum().abs() < 1e-3: 285 | raise RuntimeError('The weights must not sum to 0.') 286 | weights /= weights.sum().abs() 287 | 288 | 289 | def do_run(): 290 | 291 | if args.seed >= 0: 292 | torch.manual_seed(args.seed) 293 | 294 | kwargs = { 295 | "super_res_embed": im_tensor*0.18215, 296 | } 297 | 298 | if args.batch_size > 1: 299 | kwargs['super_res_embed'] = torch.cat(args.batch_size*[kwargs['super_res_embed']], dim=0) 300 | 301 | cur_t = None 302 | 303 | if args.ddpm: 304 | sample_fn = diffusion.ddpm_sample_loop_progressive 305 | elif args.ddim: 306 | sample_fn = diffusion.ddim_sample_loop_progressive 307 | else: 308 | sample_fn = diffusion.plms_sample_loop_progressive 309 | 310 | def save_sample(i, sample, clip_score=False): 311 | for k, image in enumerate(sample['pred_xstart']): 312 | out = TF.to_pil_image(image.add(1).div(2).clamp(0, 1)) 313 | 314 | filename = f'output/{args.prefix}{i * args.batch_size + k:05}.png' 315 | out.save(filename) 316 | 317 | def cond_fn_clip(x, t, super_res_embed=None): 318 | 319 | with torch.enable_grad(): 320 | x = x.detach().requires_grad_() 321 | n = x.shape[0] 322 | 323 | my_t = torch.ones([n], device=device, dtype=torch.long) * cur_t 324 | 325 | out = diffusion.p_mean_variance(model, x, my_t, clip_denoised=False, model_kwargs={'super_res_embed': super_res_embed}) 326 | fac = diffusion.sqrt_one_minus_alphas_cumprod[cur_t] 327 | x_in = out['pred_xstart'] * fac + x * (1 - fac) 328 | clip_in = normalize(make_cutouts(x_in.add(1).div(2))) 329 | clip_embeds = clip_model.encode_image(clip_in).float() 330 | dists = spherical_dist_loss(clip_embeds.unsqueeze(1), target_embeds.unsqueeze(0)) 331 | dists = dists.view([args.cutn, n, -1]) 332 | losses = dists.mul(weights).sum(2).mean(0) 333 | tv_losses = tv_loss(x_in) 334 | range_losses = range_loss(out['pred_xstart']) 335 | loss = losses.sum() * args.clip_guidance_scale + tv_losses.sum() * args.tv_scale + range_losses.sum() * args.range_scale 336 | return -torch.autograd.grad(loss, x)[0] 337 | 338 | def cond_fn_super(x, t, super_res_embed=None): 339 | with torch.enable_grad(): 340 | x = x.detach().requires_grad_() 341 | 342 | n = x.shape[0] 343 | 344 | my_t = torch.ones([n], device=device, dtype=torch.long) * cur_t 345 | 346 | out = diffusion.p_mean_variance(model, x, my_t, clip_denoised=False, model_kwargs={'super_res_embed': super_res_embed}) 347 | 348 | fac = diffusion.sqrt_one_minus_alphas_cumprod[cur_t] 349 | x_in = out['pred_xstart'] * fac + x * (1 - fac) 350 | 351 | x_small = F.interpolate(x_in, size=(x.shape[3]//2, x.shape[2]//2), mode='bilinear') 352 | diff = x_small - input_image_small 353 | 354 | loss = 100*(diff**2).mean([1, 2, 3]).sum() 355 | 356 | return -torch.autograd.grad(loss, x)[0] 357 | 358 | if args.init_image: 359 | init = Image.open(args.init_image).convert('RGB') 360 | init = init.resize((int(args.width), int(args.height)), Image.LANCZOS) 361 | init = TF.to_tensor(init).to(device).unsqueeze(0) 362 | init = 2*init - 1 363 | elif args.skip_timesteps > 0: 364 | init = F.interpolate(input_image, size=(args.height, args.width), mode='bicubic') 365 | else: 366 | init = None 367 | 368 | if args.text: 369 | cond_fn = cond_fn_clip 370 | elif args.guide: 371 | cond_fn = cond_fn_super 372 | else: 373 | cond_fn = None 374 | 375 | if args.avg: 376 | kwargs['super_res_embed'] = torch.cat([kwargs['super_res_embed'], kwargs['super_res_embed']], dim=0) 377 | kwargs['avg'] = True 378 | if init is not None: 379 | init = torch.cat([init, init], dim=0) 380 | 381 | for i in range(args.num_batches): 382 | cur_t = diffusion.num_timesteps - 1 383 | 384 | samples = sample_fn( 385 | model, 386 | (args.batch_size*2 if args.avg else args.batch_size, 3, args.height, args.width), 387 | clip_denoised=False, 388 | model_kwargs=kwargs, 389 | cond_fn=cond_fn, 390 | device=device, 391 | progress=True, 392 | init_image=init, 393 | skip_timesteps=args.skip_timesteps, 394 | ) 395 | 396 | for j, sample in enumerate(samples): 397 | cur_t -= 1 398 | if j % 5 == 0 and j != diffusion.num_timesteps - 1: 399 | save_sample(i, sample) 400 | 401 | save_sample(i, sample, args.clip_score) 402 | 403 | gc.collect() 404 | do_run() 405 | -------------------------------------------------------------------------------- /train.sh: -------------------------------------------------------------------------------- 1 | MODEL_FLAGS="--actual_image_size 512 --lr_warmup_steps 10000 --ema_rate 0.9999 --attention_resolutions 64,32,16 --class_cond False --diffusion_steps 1000 --image_size 64 --learn_sigma False --noise_schedule linear --num_channels 320 --num_heads 8 --num_res_blocks 2 --resblock_updown False --use_fp16 True --use_scale_shift_norm False " 2 | TRAIN_FLAGS="--lr 1e-5 --batch_size 32 --microbatch 8 --log_interval 10 --save_interval 5000 --kl_model kl.pt --resume_checkpoint diffusion.pt" 3 | export OPENAI_LOGDIR=./logs/ 4 | python scripts/image_train_stable.py --data_dir /path/to/image-and-text-files $MODEL_FLAGS $TRAIN_FLAGS 5 | -------------------------------------------------------------------------------- /train_classifier.sh: -------------------------------------------------------------------------------- 1 | MODEL_FLAGS="--actual_image_size 512 --lr_warmup_steps 0 --weight_decay 0.15 --classifier_attention_resolutions 64,32,16,8 --image_size 64 --classifier_width 128 --classifier_depth 4 --classifier_use_fp16 True " 2 | TRAIN_FLAGS="--lr 2e-5 --batch_size 20 --log_interval 10 --save_interval 10000 --kl_model kl.pt" 3 | export OPENAI_LOGDIR=./logs_classifier/ 4 | 5 | mpiexec -n 4 python scripts/classifier_train_stable.py --good_data_dir /your-images/ --bad_data_dir /laion-images/ $MODEL_FLAGS $TRAIN_FLAGS 6 | -------------------------------------------------------------------------------- /train_inpaint.sh: -------------------------------------------------------------------------------- 1 | MODEL_FLAGS="--actual_image_size 512 --lr_warmup_steps 10000 --ema_rate 0.9999 --attention_resolutions 64,32,16 --class_cond False --diffusion_steps 1000 --image_size 64 --learn_sigma False --noise_schedule linear --num_channels 320 --num_heads 8 --num_res_blocks 2 --resblock_updown False --use_fp16 True --use_scale_shift_norm False " 2 | TRAIN_FLAGS="--lr 1e-5 --batch_size 32 --microbatch 8 --log_interval 10 --save_interval 5000 --kl_model kl.pt --resume_checkpoint diffusion.pt" 3 | export OPENAI_LOGDIR=./logs/ 4 | python scripts/image_train_inpaint.py --data_dir /path/to/image-and-text-files $MODEL_FLAGS $TRAIN_FLAGS 5 | --------------------------------------------------------------------------------