abhishek-kumar commited on
Commit
ad88530
·
1 Parent(s): 30e3682

Upload 25 files

Browse files

Add improved_diffusion github files

.gitignore ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ .DS_Store
2
+ __pycache__/
3
+
LICENSE ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 CHANGED
@@ -1,4 +1,3 @@
1
- ---
2
  title: Improved Diffusion
3
  emoji: 😻
4
  colorFrom: red
@@ -7,6 +6,150 @@ sdk: gradio
7
  sdk_version: 3.23.0
8
  app_file: app.py
9
  pinned: false
10
- ---
11
 
12
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  title: Improved Diffusion
2
  emoji: 😻
3
  colorFrom: red
 
6
  sdk_version: 3.23.0
7
  app_file: app.py
8
  pinned: false
 
9
 
10
+ # improved-diffusion
11
+
12
+ This is the codebase for [Improved Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2102.09672).
13
+
14
+ # Usage
15
+
16
+ This section of the README walks through how to train and sample from a model.
17
+
18
+ ## Installation
19
+
20
+ Clone this repository and navigate to it in your terminal. Then run:
21
+
22
+ ```
23
+ pip install -e .
24
+ ```
25
+
26
+ This should install the `improved_diffusion` python package that the scripts depend on.
27
+
28
+ ## Preparing Data
29
+
30
+ The training code reads images from a directory of image files. In the [datasets](datasets) folder, we have provided instructions/scripts for preparing these directories for ImageNet, LSUN bedrooms, and CIFAR-10.
31
+
32
+ For creating your own dataset, simply dump all of your images into a directory with ".jpg", ".jpeg", or ".png" extensions. If you wish to train a class-conditional model, name the files like "mylabel1_XXX.jpg", "mylabel2_YYY.jpg", etc., so that the data loader knows that "mylabel1" and "mylabel2" are the labels. Subdirectories will automatically be enumerated as well, so the images can be organized into a recursive structure (although the directory names will be ignored, and the underscore prefixes are used as names).
33
+
34
+ The images will automatically be scaled and center-cropped by the data-loading pipeline. Simply pass `--data_dir path/to/images` to the training script, and it will take care of the rest.
35
+
36
+ ## Training
37
+
38
+ To train your model, you should first decide some hyperparameters. We will split up our hyperparameters into three groups: model architecture, diffusion process, and training flags. Here are some reasonable defaults for a baseline:
39
+
40
+ ```
41
+ MODEL_FLAGS="--image_size 64 --num_channels 128 --num_res_blocks 3"
42
+ DIFFUSION_FLAGS="--diffusion_steps 4000 --noise_schedule linear"
43
+ TRAIN_FLAGS="--lr 1e-4 --batch_size 128"
44
+ ```
45
+
46
+ Here are some changes we experiment with, and how to set them in the flags:
47
+
48
+ * **Learned sigmas:** add `--learn_sigma True` to `MODEL_FLAGS`
49
+ * **Cosine schedule:** change `--noise_schedule linear` to `--noise_schedule cosine`
50
+ * **Importance-sampled VLB:** add `--use_kl True` to `DIFFUSION_FLAGS` and add `--schedule_sampler loss-second-moment` to `TRAIN_FLAGS`.
51
+ * **Class-conditional:** add `--class_cond True` to `MODEL_FLAGS`.
52
+
53
+ Once you have setup your hyper-parameters, you can run an experiment like so:
54
+
55
+ ```
56
+ python scripts/image_train.py --data_dir path/to/images $MODEL_FLAGS $DIFFUSION_FLAGS $TRAIN_FLAGS
57
+ ```
58
+
59
+ You may also want to train in a distributed manner. In this case, run the same command with `mpiexec`:
60
+
61
+ ```
62
+ mpiexec -n $NUM_GPUS python scripts/image_train.py --data_dir path/to/images $MODEL_FLAGS $DIFFUSION_FLAGS $TRAIN_FLAGS
63
+ ```
64
+
65
+ When training in a distributed manner, you must manually divide the `--batch_size` argument by the number of ranks. In lieu of distributed training, you may use `--microbatch 16` (or `--microbatch 1` in extreme memory-limited cases) to reduce memory usage.
66
+
67
+ The logs and saved models will be written to a logging directory determined by the `OPENAI_LOGDIR` environment variable. If it is not set, then a temporary directory will be created in `/tmp`.
68
+
69
+ ## Sampling
70
+
71
+ The above training script saves checkpoints to `.pt` files in the logging directory. These checkpoints will have names like `ema_0.9999_200000.pt` and `model200000.pt`. You will likely want to sample from the EMA models, since those produce much better samples.
72
+
73
+ Once you have a path to your model, you can generate a large batch of samples like so:
74
+
75
+ ```
76
+ python scripts/image_sample.py --model_path /path/to/model.pt $MODEL_FLAGS $DIFFUSION_FLAGS
77
+ ```
78
+
79
+ Again, this will save results to a logging directory. Samples are saved as a large `npz` file, where `arr_0` in the file is a large batch of samples.
80
+
81
+ Just like for training, you can run `image_sample.py` through MPI to use multiple GPUs and machines.
82
+
83
+ You can change the number of sampling steps using the `--timestep_respacing` argument. For example, `--timestep_respacing 250` uses 250 steps to sample. Passing `--timestep_respacing ddim250` is similar, but uses the uniform stride from the [DDIM paper](https://arxiv.org/abs/2010.02502) rather than our stride.
84
+
85
+ To sample using [DDIM](https://arxiv.org/abs/2010.02502), pass `--use_ddim True`.
86
+
87
+ ## Models and Hyperparameters
88
+
89
+ This section includes model checkpoints and run flags for the main models in the paper.
90
+
91
+ Note that the batch sizes are specified for single-GPU training, even though most of these runs will not naturally fit on a single GPU. To address this, either set `--microbatch` to a small value (e.g. 4) to train on one GPU, or run with MPI and divide `--batch_size` by the number of GPUs.
92
+
93
+ Unconditional ImageNet-64 with our `L_hybrid` objective and cosine noise schedule [[checkpoint](https://openaipublic.blob.core.windows.net/diffusion/march-2021/imagenet64_uncond_100M_1500K.pt)]:
94
+
95
+ ```bash
96
+ MODEL_FLAGS="--image_size 64 --num_channels 128 --num_res_blocks 3 --learn_sigma True"
97
+ DIFFUSION_FLAGS="--diffusion_steps 4000 --noise_schedule cosine"
98
+ TRAIN_FLAGS="--lr 1e-4 --batch_size 128"
99
+ ```
100
+
101
+ Unconditional CIFAR-10 with our `L_hybrid` objective and cosine noise schedule [[checkpoint](https://openaipublic.blob.core.windows.net/diffusion/march-2021/cifar10_uncond_50M_500K.pt)]:
102
+
103
+ ```bash
104
+ MODEL_FLAGS="--image_size 32 --num_channels 128 --num_res_blocks 3 --learn_sigma True --dropout 0.3"
105
+ DIFFUSION_FLAGS="--diffusion_steps 4000 --noise_schedule cosine"
106
+ TRAIN_FLAGS="--lr 1e-4 --batch_size 128"
107
+ ```
108
+
109
+ Class-conditional ImageNet-64 model (270M parameters, trained for 250K iterations) [[checkpoint](https://openaipublic.blob.core.windows.net/diffusion/march-2021/imagenet64_cond_270M_250K.pt)]:
110
+
111
+ ```bash
112
+ MODEL_FLAGS="--image_size 64 --num_channels 192 --num_res_blocks 3 --learn_sigma True --class_cond True"
113
+ DIFFUSION_FLAGS="--diffusion_steps 4000 --noise_schedule cosine --rescale_learned_sigmas False --rescale_timesteps False"
114
+ TRAIN_FLAGS="--lr 3e-4 --batch_size 2048"
115
+ ```
116
+
117
+ Upsampling 256x256 model (280M parameters, trained for 500K iterations) [[checkpoint](https://openaipublic.blob.core.windows.net/diffusion/march-2021/upsample_cond_500K.pt)]:
118
+
119
+ ```bash
120
+ MODEL_FLAGS="--num_channels 192 --num_res_blocks 2 --learn_sigma True --class_cond True"
121
+ DIFFUSION_FLAGS="--diffusion_steps 4000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False"
122
+ TRAIN_FLAGS="--lr 3e-4 --batch_size 256"
123
+ ```
124
+
125
+ LSUN bedroom model (lr=1e-4) [[checkpoint](https://openaipublic.blob.core.windows.net/diffusion/march-2021/lsun_uncond_100M_1200K_bs128.pt)]:
126
+
127
+ ```bash
128
+ MODEL_FLAGS="--image_size 256 --num_channels 128 --num_res_blocks 2 --num_heads 1 --learn_sigma True --use_scale_shift_norm False --attention_resolutions 16"
129
+ DIFFUSION_FLAGS="--diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False"
130
+ TRAIN_FLAGS="--lr 1e-4 --batch_size 128"
131
+ ```
132
+
133
+ LSUN bedroom model (lr=2e-5) [[checkpoint](https://openaipublic.blob.core.windows.net/diffusion/march-2021/lsun_uncond_100M_2400K_bs64.pt)]:
134
+
135
+ ```bash
136
+ MODEL_FLAGS="--image_size 256 --num_channels 128 --num_res_blocks 2 --num_heads 1 --learn_sigma True --use_scale_shift_norm False --attention_resolutions 16"
137
+ DIFFUSION_FLAGS="--diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False --use_scale_shift_norm False"
138
+ TRAIN_FLAGS="--lr 2e-5 --batch_size 128"
139
+ ```
140
+
141
+ Unconditional ImageNet-64 with the `L_vlb` objective and cosine noise schedule [[checkpoint](https://openaipublic.blob.core.windows.net/diffusion/march-2021/imagenet64_uncond_vlb_100M_1500K.pt)]:
142
+
143
+ ```bash
144
+ MODEL_FLAGS="--image_size 64 --num_channels 128 --num_res_blocks 3 --learn_sigma True"
145
+ DIFFUSION_FLAGS="--diffusion_steps 4000 --noise_schedule cosine --use_kl True"
146
+ TRAIN_FLAGS="--lr 1e-4 --batch_size 128 --schedule_sampler loss-second-moment"
147
+ ```
148
+
149
+ Unconditional CIFAR-10 with the `L_vlb` objective and cosine noise schedule [[checkpoint](https://openaipublic.blob.core.windows.net/diffusion/march-2021/cifar10_uncond_vlb_50M_500K.pt)]:
150
+
151
+ ```bash
152
+ MODEL_FLAGS="--image_size 32 --num_channels 128 --num_res_blocks 3 --learn_sigma True --dropout 0.3"
153
+ DIFFUSION_FLAGS="--diffusion_steps 4000 --noise_schedule cosine --use_kl True"
154
+ TRAIN_FLAGS="--lr 1e-4 --batch_size 128 --schedule_sampler loss-second-moment"
155
+ ```
datasets/README.md ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Downloading datasets
2
+
3
+ This directory includes instructions and scripts for downloading ImageNet, LSUN bedrooms, and CIFAR-10 for use in this codebase.
4
+
5
+ ## ImageNet-64
6
+
7
+ To download unconditional ImageNet-64, go to [this page on image-net.org](http://www.image-net.org/small/download.php) and click on "Train (64x64)". Simply download the file and unzip it, and use the resulting directory as the data directory (the `--data_dir` argument for the training script).
8
+
9
+ ## Class-conditional ImageNet
10
+
11
+ 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.
12
+
13
+ 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:
14
+
15
+ ```
16
+ for file in *.tar; do tar xf "$file"; rm "$file"; done
17
+ ```
18
+
19
+ This will extract and remove each tar file in turn.
20
+
21
+ 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.
22
+
23
+ ## CIFAR-10
24
+
25
+ For CIFAR-10, we created a script [cifar10.py](cifar10.py) that creates `cifar_train` and `cifar_test` directories. These directories contain files named like `truck_49997.png`, so that the class name is discernable to the data loader.
26
+
27
+ The `cifar_train` and `cifar_test` directories can be passed directly to the training scripts via the `--data_dir` argument.
28
+
29
+ ## LSUN bedroom
30
+
31
+ 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:
32
+
33
+ ```
34
+ python lsun_bedroom.py bedroom_train_lmdb lsun_train_output_dir
35
+ ```
36
+
37
+ This creates a directory called `lsun_train_output_dir`. This directory can be passed to the training scripts via the `--data_dir` argument.
datasets/cifar10.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import tempfile
3
+
4
+ import torchvision
5
+ from tqdm.auto import tqdm
6
+
7
+ CLASSES = (
8
+ "plane",
9
+ "car",
10
+ "bird",
11
+ "cat",
12
+ "deer",
13
+ "dog",
14
+ "frog",
15
+ "horse",
16
+ "ship",
17
+ "truck",
18
+ )
19
+
20
+
21
+ def main():
22
+ for split in ["train", "test"]:
23
+ out_dir = f"cifar_{split}"
24
+ if os.path.exists(out_dir):
25
+ print(f"skipping split {split} since {out_dir} already exists.")
26
+ continue
27
+
28
+ print("downloading...")
29
+ with tempfile.TemporaryDirectory() as tmp_dir:
30
+ dataset = torchvision.datasets.CIFAR10(
31
+ root=tmp_dir, train=split == "train", download=True
32
+ )
33
+
34
+ print("dumping images...")
35
+ os.mkdir(out_dir)
36
+ for i in tqdm(range(len(dataset))):
37
+ image, label = dataset[i]
38
+ filename = os.path.join(out_dir, f"{CLASSES[label]}_{i:05d}.png")
39
+ image.save(filename)
40
+
41
+
42
+ if __name__ == "__main__":
43
+ main()
datasets/lsun_bedroom.py ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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()
improved_diffusion/__init__.py ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ """
2
+ Codebase for "Improved Denoising Diffusion Probabilistic Models".
3
+ """
improved_diffusion/dist_util.py ADDED
@@ -0,0 +1,82 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+
28
+ comm = MPI.COMM_WORLD
29
+ backend = "gloo" if not th.cuda.is_available() else "nccl"
30
+
31
+ if backend == "gloo":
32
+ hostname = "localhost"
33
+ else:
34
+ hostname = socket.gethostbyname(socket.getfqdn())
35
+ os.environ["MASTER_ADDR"] = comm.bcast(hostname, root=0)
36
+ os.environ["RANK"] = str(comm.rank)
37
+ os.environ["WORLD_SIZE"] = str(comm.size)
38
+
39
+ port = comm.bcast(_find_free_port(), root=0)
40
+ os.environ["MASTER_PORT"] = str(port)
41
+ dist.init_process_group(backend=backend, init_method="env://")
42
+
43
+
44
+ def dev():
45
+ """
46
+ Get the device to use for torch.distributed.
47
+ """
48
+ if th.cuda.is_available():
49
+ return th.device(f"cuda:{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}")
50
+ return th.device("cpu")
51
+
52
+
53
+ def load_state_dict(path, **kwargs):
54
+ """
55
+ Load a PyTorch file without redundant fetches across MPI ranks.
56
+ """
57
+ if MPI.COMM_WORLD.Get_rank() == 0:
58
+ with bf.BlobFile(path, "rb") as f:
59
+ data = f.read()
60
+ else:
61
+ data = None
62
+ data = MPI.COMM_WORLD.bcast(data)
63
+ return th.load(io.BytesIO(data), **kwargs)
64
+
65
+
66
+ def sync_params(params):
67
+ """
68
+ Synchronize a sequence of Tensors across ranks from rank 0.
69
+ """
70
+ for p in params:
71
+ with th.no_grad():
72
+ dist.broadcast(p, 0)
73
+
74
+
75
+ def _find_free_port():
76
+ try:
77
+ s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
78
+ s.bind(("", 0))
79
+ s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
80
+ return s.getsockname()[1]
81
+ finally:
82
+ s.close()
improved_diffusion/fp16_util.py ADDED
@@ -0,0 +1,76 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Helpers to train with 16-bit precision.
3
+ """
4
+
5
+ import torch.nn as nn
6
+ from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
7
+
8
+
9
+ def convert_module_to_f16(l):
10
+ """
11
+ Convert primitive modules to float16.
12
+ """
13
+ if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
14
+ l.weight.data = l.weight.data.half()
15
+ l.bias.data = l.bias.data.half()
16
+
17
+
18
+ def convert_module_to_f32(l):
19
+ """
20
+ Convert primitive modules to float32, undoing convert_module_to_f16().
21
+ """
22
+ if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
23
+ l.weight.data = l.weight.data.float()
24
+ l.bias.data = l.bias.data.float()
25
+
26
+
27
+ def make_master_params(model_params):
28
+ """
29
+ Copy model parameters into a (differently-shaped) list of full-precision
30
+ parameters.
31
+ """
32
+ master_params = _flatten_dense_tensors(
33
+ [param.detach().float() for param in model_params]
34
+ )
35
+ master_params = nn.Parameter(master_params)
36
+ master_params.requires_grad = True
37
+ return [master_params]
38
+
39
+
40
+ def model_grads_to_master_grads(model_params, master_params):
41
+ """
42
+ Copy the gradients from the model parameters into the master parameters
43
+ from make_master_params().
44
+ """
45
+ master_params[0].grad = _flatten_dense_tensors(
46
+ [param.grad.data.detach().float() for param in model_params]
47
+ )
48
+
49
+
50
+ def master_params_to_model_params(model_params, master_params):
51
+ """
52
+ Copy the master parameter data back into the model parameters.
53
+ """
54
+ # Without copying to a list, if a generator is passed, this will
55
+ # silently not copy any parameters.
56
+ model_params = list(model_params)
57
+
58
+ for param, master_param in zip(
59
+ model_params, unflatten_master_params(model_params, master_params)
60
+ ):
61
+ param.detach().copy_(master_param)
62
+
63
+
64
+ def unflatten_master_params(model_params, master_params):
65
+ """
66
+ Unflatten the master parameters to look like model_params.
67
+ """
68
+ return _unflatten_dense_tensors(master_params[0].detach(), model_params)
69
+
70
+
71
+ def zero_grad(model_params):
72
+ for param in model_params:
73
+ # Taken from https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.add_param_group
74
+ if param.grad is not None:
75
+ param.grad.detach_()
76
+ param.grad.zero_()
improved_diffusion/gaussian_diffusion.py ADDED
@@ -0,0 +1,841 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ This code started out as a PyTorch port of Ho et al's diffusion models:
3
+ https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py
4
+
5
+ Docstrings have been added, as well as DDIM sampling and a new collection of beta schedules.
6
+ """
7
+
8
+ import enum
9
+ import math
10
+
11
+ import numpy as np
12
+ import torch as th
13
+
14
+ from .nn import mean_flat
15
+ from .losses import normal_kl, discretized_gaussian_log_likelihood
16
+
17
+
18
+ def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
19
+ """
20
+ Get a pre-defined beta schedule for the given name.
21
+
22
+ The beta schedule library consists of beta schedules which remain similar
23
+ in the limit of num_diffusion_timesteps.
24
+ Beta schedules may be added, but should not be removed or changed once
25
+ they are committed to maintain backwards compatibility.
26
+ """
27
+ if schedule_name == "linear":
28
+ # Linear schedule from Ho et al, extended to work for any number of
29
+ # diffusion steps.
30
+ scale = 1000 / num_diffusion_timesteps
31
+ beta_start = scale * 0.0001
32
+ beta_end = scale * 0.02
33
+ return np.linspace(
34
+ beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64
35
+ )
36
+ elif schedule_name == "cosine":
37
+ return betas_for_alpha_bar(
38
+ num_diffusion_timesteps,
39
+ lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
40
+ )
41
+ else:
42
+ raise NotImplementedError(f"unknown beta schedule: {schedule_name}")
43
+
44
+
45
+ def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
46
+ """
47
+ Create a beta schedule that discretizes the given alpha_t_bar function,
48
+ which defines the cumulative product of (1-beta) over time from t = [0,1].
49
+
50
+ :param num_diffusion_timesteps: the number of betas to produce.
51
+ :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
52
+ produces the cumulative product of (1-beta) up to that
53
+ part of the diffusion process.
54
+ :param max_beta: the maximum beta to use; use values lower than 1 to
55
+ prevent singularities.
56
+ """
57
+ betas = []
58
+ for i in range(num_diffusion_timesteps):
59
+ t1 = i / num_diffusion_timesteps
60
+ t2 = (i + 1) / num_diffusion_timesteps
61
+ betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
62
+ return np.array(betas)
63
+
64
+
65
+ class ModelMeanType(enum.Enum):
66
+ """
67
+ Which type of output the model predicts.
68
+ """
69
+
70
+ PREVIOUS_X = enum.auto() # the model predicts x_{t-1}
71
+ START_X = enum.auto() # the model predicts x_0
72
+ EPSILON = enum.auto() # the model predicts epsilon
73
+
74
+
75
+ class ModelVarType(enum.Enum):
76
+ """
77
+ What is used as the model's output variance.
78
+
79
+ The LEARNED_RANGE option has been added to allow the model to predict
80
+ values between FIXED_SMALL and FIXED_LARGE, making its job easier.
81
+ """
82
+
83
+ LEARNED = enum.auto()
84
+ FIXED_SMALL = enum.auto()
85
+ FIXED_LARGE = enum.auto()
86
+ LEARNED_RANGE = enum.auto()
87
+
88
+
89
+ class LossType(enum.Enum):
90
+ MSE = enum.auto() # use raw MSE loss (and KL when learning variances)
91
+ RESCALED_MSE = (
92
+ enum.auto()
93
+ ) # use raw MSE loss (with RESCALED_KL when learning variances)
94
+ KL = enum.auto() # use the variational lower-bound
95
+ RESCALED_KL = enum.auto() # like KL, but rescale to estimate the full VLB
96
+
97
+ def is_vb(self):
98
+ return self == LossType.KL or self == LossType.RESCALED_KL
99
+
100
+
101
+ class GaussianDiffusion:
102
+ """
103
+ Utilities for training and sampling diffusion models.
104
+
105
+ Ported directly from here, and then adapted over time to further experimentation.
106
+ https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42
107
+
108
+ :param betas: a 1-D numpy array of betas for each diffusion timestep,
109
+ starting at T and going to 1.
110
+ :param model_mean_type: a ModelMeanType determining what the model outputs.
111
+ :param model_var_type: a ModelVarType determining how variance is output.
112
+ :param loss_type: a LossType determining the loss function to use.
113
+ :param rescale_timesteps: if True, pass floating point timesteps into the
114
+ model so that they are always scaled like in the
115
+ original paper (0 to 1000).
116
+ """
117
+
118
+ def __init__(
119
+ self,
120
+ *,
121
+ betas,
122
+ model_mean_type,
123
+ model_var_type,
124
+ loss_type,
125
+ rescale_timesteps=False,
126
+ ):
127
+ self.model_mean_type = model_mean_type
128
+ self.model_var_type = model_var_type
129
+ self.loss_type = loss_type
130
+ self.rescale_timesteps = rescale_timesteps
131
+
132
+ # Use float64 for accuracy.
133
+ betas = np.array(betas, dtype=np.float64)
134
+ self.betas = betas
135
+ assert len(betas.shape) == 1, "betas must be 1-D"
136
+ assert (betas > 0).all() and (betas <= 1).all()
137
+
138
+ self.num_timesteps = int(betas.shape[0])
139
+
140
+ alphas = 1.0 - betas
141
+ self.alphas_cumprod = np.cumprod(alphas, axis=0)
142
+ self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
143
+ self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
144
+ assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)
145
+
146
+ # calculations for diffusion q(x_t | x_{t-1}) and others
147
+ self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
148
+ self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
149
+ self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
150
+ self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
151
+ self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)
152
+
153
+ # calculations for posterior q(x_{t-1} | x_t, x_0)
154
+ self.posterior_variance = (
155
+ betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
156
+ )
157
+ # log calculation clipped because the posterior variance is 0 at the
158
+ # beginning of the diffusion chain.
159
+ self.posterior_log_variance_clipped = np.log(
160
+ np.append(self.posterior_variance[1], self.posterior_variance[1:])
161
+ )
162
+ self.posterior_mean_coef1 = (
163
+ betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
164
+ )
165
+ self.posterior_mean_coef2 = (
166
+ (1.0 - self.alphas_cumprod_prev)
167
+ * np.sqrt(alphas)
168
+ / (1.0 - self.alphas_cumprod)
169
+ )
170
+
171
+ def q_mean_variance(self, x_start, t):
172
+ """
173
+ Get the distribution q(x_t | x_0).
174
+
175
+ :param x_start: the [N x C x ...] tensor of noiseless inputs.
176
+ :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
177
+ :return: A tuple (mean, variance, log_variance), all of x_start's shape.
178
+ """
179
+ mean = (
180
+ _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
181
+ )
182
+ variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
183
+ log_variance = _extract_into_tensor(
184
+ self.log_one_minus_alphas_cumprod, t, x_start.shape
185
+ )
186
+ return mean, variance, log_variance
187
+
188
+ def q_sample(self, x_start, t, noise=None):
189
+ """
190
+ Diffuse the data for a given number of diffusion steps.
191
+
192
+ In other words, sample from q(x_t | x_0).
193
+
194
+ :param x_start: the initial data batch.
195
+ :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
196
+ :param noise: if specified, the split-out normal noise.
197
+ :return: A noisy version of x_start.
198
+ """
199
+ if noise is None:
200
+ noise = th.randn_like(x_start)
201
+ assert noise.shape == x_start.shape
202
+ return (
203
+ _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
204
+ + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
205
+ * noise
206
+ )
207
+
208
+ def q_posterior_mean_variance(self, x_start, x_t, t):
209
+ """
210
+ Compute the mean and variance of the diffusion posterior:
211
+
212
+ q(x_{t-1} | x_t, x_0)
213
+
214
+ """
215
+ assert x_start.shape == x_t.shape
216
+ posterior_mean = (
217
+ _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start
218
+ + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
219
+ )
220
+ posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape)
221
+ posterior_log_variance_clipped = _extract_into_tensor(
222
+ self.posterior_log_variance_clipped, t, x_t.shape
223
+ )
224
+ assert (
225
+ posterior_mean.shape[0]
226
+ == posterior_variance.shape[0]
227
+ == posterior_log_variance_clipped.shape[0]
228
+ == x_start.shape[0]
229
+ )
230
+ return posterior_mean, posterior_variance, posterior_log_variance_clipped
231
+
232
+ def p_mean_variance(
233
+ self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None
234
+ ):
235
+ """
236
+ Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
237
+ the initial x, x_0.
238
+
239
+ :param model: the model, which takes a signal and a batch of timesteps
240
+ as input.
241
+ :param x: the [N x C x ...] tensor at time t.
242
+ :param t: a 1-D Tensor of timesteps.
243
+ :param clip_denoised: if True, clip the denoised signal into [-1, 1].
244
+ :param denoised_fn: if not None, a function which applies to the
245
+ x_start prediction before it is used to sample. Applies before
246
+ clip_denoised.
247
+ :param model_kwargs: if not None, a dict of extra keyword arguments to
248
+ pass to the model. This can be used for conditioning.
249
+ :return: a dict with the following keys:
250
+ - 'mean': the model mean output.
251
+ - 'variance': the model variance output.
252
+ - 'log_variance': the log of 'variance'.
253
+ - 'pred_xstart': the prediction for x_0.
254
+ """
255
+ if model_kwargs is None:
256
+ model_kwargs = {}
257
+
258
+ B, C = x.shape[:2]
259
+ assert t.shape == (B,)
260
+ model_output = model(x, self._scale_timesteps(t), **model_kwargs)
261
+
262
+ if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:
263
+ assert model_output.shape == (B, C * 2, *x.shape[2:])
264
+ model_output, model_var_values = th.split(model_output, C, dim=1)
265
+ if self.model_var_type == ModelVarType.LEARNED:
266
+ model_log_variance = model_var_values
267
+ model_variance = th.exp(model_log_variance)
268
+ else:
269
+ min_log = _extract_into_tensor(
270
+ self.posterior_log_variance_clipped, t, x.shape
271
+ )
272
+ max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)
273
+ # The model_var_values is [-1, 1] for [min_var, max_var].
274
+ frac = (model_var_values + 1) / 2
275
+ model_log_variance = frac * max_log + (1 - frac) * min_log
276
+ model_variance = th.exp(model_log_variance)
277
+ else:
278
+ model_variance, model_log_variance = {
279
+ # for fixedlarge, we set the initial (log-)variance like so
280
+ # to get a better decoder log likelihood.
281
+ ModelVarType.FIXED_LARGE: (
282
+ np.append(self.posterior_variance[1], self.betas[1:]),
283
+ np.log(np.append(self.posterior_variance[1], self.betas[1:])),
284
+ ),
285
+ ModelVarType.FIXED_SMALL: (
286
+ self.posterior_variance,
287
+ self.posterior_log_variance_clipped,
288
+ ),
289
+ }[self.model_var_type]
290
+ model_variance = _extract_into_tensor(model_variance, t, x.shape)
291
+ model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape)
292
+
293
+ def process_xstart(x):
294
+ if denoised_fn is not None:
295
+ x = denoised_fn(x)
296
+ if clip_denoised:
297
+ return x.clamp(-1, 1)
298
+ return x
299
+
300
+ if self.model_mean_type == ModelMeanType.PREVIOUS_X:
301
+ pred_xstart = process_xstart(
302
+ self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output)
303
+ )
304
+ model_mean = model_output
305
+ elif self.model_mean_type in [ModelMeanType.START_X, ModelMeanType.EPSILON]:
306
+ if self.model_mean_type == ModelMeanType.START_X:
307
+ pred_xstart = process_xstart(model_output)
308
+ else:
309
+ pred_xstart = process_xstart(
310
+ self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)
311
+ )
312
+ model_mean, _, _ = self.q_posterior_mean_variance(
313
+ x_start=pred_xstart, x_t=x, t=t
314
+ )
315
+ else:
316
+ raise NotImplementedError(self.model_mean_type)
317
+
318
+ assert (
319
+ model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
320
+ )
321
+ return {
322
+ "mean": model_mean,
323
+ "variance": model_variance,
324
+ "log_variance": model_log_variance,
325
+ "pred_xstart": pred_xstart,
326
+ }
327
+
328
+ def _predict_xstart_from_eps(self, x_t, t, eps):
329
+ assert x_t.shape == eps.shape
330
+ return (
331
+ _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
332
+ - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps
333
+ )
334
+
335
+ def _predict_xstart_from_xprev(self, x_t, t, xprev):
336
+ assert x_t.shape == xprev.shape
337
+ return ( # (xprev - coef2*x_t) / coef1
338
+ _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev
339
+ - _extract_into_tensor(
340
+ self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape
341
+ )
342
+ * x_t
343
+ )
344
+
345
+ def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
346
+ return (
347
+ _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
348
+ - pred_xstart
349
+ ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
350
+
351
+ def _scale_timesteps(self, t):
352
+ if self.rescale_timesteps:
353
+ return t.float() * (1000.0 / self.num_timesteps)
354
+ return t
355
+
356
+ def p_sample(
357
+ self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None
358
+ ):
359
+ """
360
+ Sample x_{t-1} from the model at the given timestep.
361
+
362
+ :param model: the model to sample from.
363
+ :param x: the current tensor at x_{t-1}.
364
+ :param t: the value of t, starting at 0 for the first diffusion step.
365
+ :param clip_denoised: if True, clip the x_start prediction to [-1, 1].
366
+ :param denoised_fn: if not None, a function which applies to the
367
+ x_start prediction before it is used to sample.
368
+ :param model_kwargs: if not None, a dict of extra keyword arguments to
369
+ pass to the model. This can be used for conditioning.
370
+ :return: a dict containing the following keys:
371
+ - 'sample': a random sample from the model.
372
+ - 'pred_xstart': a prediction of x_0.
373
+ """
374
+ out = self.p_mean_variance(
375
+ model,
376
+ x,
377
+ t,
378
+ clip_denoised=clip_denoised,
379
+ denoised_fn=denoised_fn,
380
+ model_kwargs=model_kwargs,
381
+ )
382
+ noise = th.randn_like(x)
383
+ nonzero_mask = (
384
+ (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
385
+ ) # no noise when t == 0
386
+ sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise
387
+ return {"sample": sample, "pred_xstart": out["pred_xstart"]}
388
+
389
+ def p_sample_loop(
390
+ self,
391
+ model,
392
+ shape,
393
+ noise=None,
394
+ clip_denoised=True,
395
+ denoised_fn=None,
396
+ model_kwargs=None,
397
+ device=None,
398
+ progress=False,
399
+ ):
400
+ """
401
+ Generate samples from the model.
402
+
403
+ :param model: the model module.
404
+ :param shape: the shape of the samples, (N, C, H, W).
405
+ :param noise: if specified, the noise from the encoder to sample.
406
+ Should be of the same shape as `shape`.
407
+ :param clip_denoised: if True, clip x_start predictions to [-1, 1].
408
+ :param denoised_fn: if not None, a function which applies to the
409
+ x_start prediction before it is used to sample.
410
+ :param model_kwargs: if not None, a dict of extra keyword arguments to
411
+ pass to the model. This can be used for conditioning.
412
+ :param device: if specified, the device to create the samples on.
413
+ If not specified, use a model parameter's device.
414
+ :param progress: if True, show a tqdm progress bar.
415
+ :return: a non-differentiable batch of samples.
416
+ """
417
+ final = None
418
+ for sample in self.p_sample_loop_progressive(
419
+ model,
420
+ shape,
421
+ noise=noise,
422
+ clip_denoised=clip_denoised,
423
+ denoised_fn=denoised_fn,
424
+ model_kwargs=model_kwargs,
425
+ device=device,
426
+ progress=progress,
427
+ ):
428
+ final = sample
429
+ return final["sample"]
430
+
431
+ def p_sample_loop_progressive(
432
+ self,
433
+ model,
434
+ shape,
435
+ noise=None,
436
+ clip_denoised=True,
437
+ denoised_fn=None,
438
+ model_kwargs=None,
439
+ device=None,
440
+ progress=False,
441
+ ):
442
+ """
443
+ Generate samples from the model and yield intermediate samples from
444
+ each timestep of diffusion.
445
+
446
+ Arguments are the same as p_sample_loop().
447
+ Returns a generator over dicts, where each dict is the return value of
448
+ p_sample().
449
+ """
450
+ if device is None:
451
+ device = next(model.parameters()).device
452
+ assert isinstance(shape, (tuple, list))
453
+ if noise is not None:
454
+ img = noise
455
+ else:
456
+ img = th.randn(*shape, device=device)
457
+ indices = list(range(self.num_timesteps))[::-1]
458
+
459
+ if progress:
460
+ # Lazy import so that we don't depend on tqdm.
461
+ from tqdm.auto import tqdm
462
+
463
+ indices = tqdm(indices)
464
+
465
+ for i in indices:
466
+ t = th.tensor([i] * shape[0], device=device)
467
+ with th.no_grad():
468
+ out = self.p_sample(
469
+ model,
470
+ img,
471
+ t,
472
+ clip_denoised=clip_denoised,
473
+ denoised_fn=denoised_fn,
474
+ model_kwargs=model_kwargs,
475
+ )
476
+ yield out
477
+ img = out["sample"]
478
+
479
+ def ddim_sample(
480
+ self,
481
+ model,
482
+ x,
483
+ t,
484
+ clip_denoised=True,
485
+ denoised_fn=None,
486
+ model_kwargs=None,
487
+ eta=0.0,
488
+ ):
489
+ """
490
+ Sample x_{t-1} from the model using DDIM.
491
+
492
+ Same usage as p_sample().
493
+ """
494
+ out = self.p_mean_variance(
495
+ model,
496
+ x,
497
+ t,
498
+ clip_denoised=clip_denoised,
499
+ denoised_fn=denoised_fn,
500
+ model_kwargs=model_kwargs,
501
+ )
502
+ # Usually our model outputs epsilon, but we re-derive it
503
+ # in case we used x_start or x_prev prediction.
504
+ eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"])
505
+ alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
506
+ alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
507
+ sigma = (
508
+ eta
509
+ * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))
510
+ * th.sqrt(1 - alpha_bar / alpha_bar_prev)
511
+ )
512
+ # Equation 12.
513
+ noise = th.randn_like(x)
514
+ mean_pred = (
515
+ out["pred_xstart"] * th.sqrt(alpha_bar_prev)
516
+ + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps
517
+ )
518
+ nonzero_mask = (
519
+ (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
520
+ ) # no noise when t == 0
521
+ sample = mean_pred + nonzero_mask * sigma * noise
522
+ return {"sample": sample, "pred_xstart": out["pred_xstart"]}
523
+
524
+ def ddim_reverse_sample(
525
+ self,
526
+ model,
527
+ x,
528
+ t,
529
+ clip_denoised=True,
530
+ denoised_fn=None,
531
+ model_kwargs=None,
532
+ eta=0.0,
533
+ ):
534
+ """
535
+ Sample x_{t+1} from the model using DDIM reverse ODE.
536
+ """
537
+ assert eta == 0.0, "Reverse ODE only for deterministic path"
538
+ out = self.p_mean_variance(
539
+ model,
540
+ x,
541
+ t,
542
+ clip_denoised=clip_denoised,
543
+ denoised_fn=denoised_fn,
544
+ model_kwargs=model_kwargs,
545
+ )
546
+ # Usually our model outputs epsilon, but we re-derive it
547
+ # in case we used x_start or x_prev prediction.
548
+ eps = (
549
+ _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x
550
+ - out["pred_xstart"]
551
+ ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape)
552
+ alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)
553
+
554
+ # Equation 12. reversed
555
+ mean_pred = (
556
+ out["pred_xstart"] * th.sqrt(alpha_bar_next)
557
+ + th.sqrt(1 - alpha_bar_next) * eps
558
+ )
559
+
560
+ return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]}
561
+
562
+ def ddim_sample_loop(
563
+ self,
564
+ model,
565
+ shape,
566
+ noise=None,
567
+ clip_denoised=True,
568
+ denoised_fn=None,
569
+ model_kwargs=None,
570
+ device=None,
571
+ progress=False,
572
+ eta=0.0,
573
+ ):
574
+ """
575
+ Generate samples from the model using DDIM.
576
+
577
+ Same usage as p_sample_loop().
578
+ """
579
+ final = None
580
+ for sample in self.ddim_sample_loop_progressive(
581
+ model,
582
+ shape,
583
+ noise=noise,
584
+ clip_denoised=clip_denoised,
585
+ denoised_fn=denoised_fn,
586
+ model_kwargs=model_kwargs,
587
+ device=device,
588
+ progress=progress,
589
+ eta=eta,
590
+ ):
591
+ final = sample
592
+ return final["sample"]
593
+
594
+ def ddim_sample_loop_progressive(
595
+ self,
596
+ model,
597
+ shape,
598
+ noise=None,
599
+ clip_denoised=True,
600
+ denoised_fn=None,
601
+ model_kwargs=None,
602
+ device=None,
603
+ progress=False,
604
+ eta=0.0,
605
+ ):
606
+ """
607
+ Use DDIM to sample from the model and yield intermediate samples from
608
+ each timestep of DDIM.
609
+
610
+ Same usage as p_sample_loop_progressive().
611
+ """
612
+ if device is None:
613
+ device = next(model.parameters()).device
614
+ assert isinstance(shape, (tuple, list))
615
+ if noise is not None:
616
+ img = noise
617
+ else:
618
+ img = th.randn(*shape, device=device)
619
+ indices = list(range(self.num_timesteps))[::-1]
620
+
621
+ if progress:
622
+ # Lazy import so that we don't depend on tqdm.
623
+ from tqdm.auto import tqdm
624
+
625
+ indices = tqdm(indices)
626
+
627
+ for i in indices:
628
+ t = th.tensor([i] * shape[0], device=device)
629
+ with th.no_grad():
630
+ out = self.ddim_sample(
631
+ model,
632
+ img,
633
+ t,
634
+ clip_denoised=clip_denoised,
635
+ denoised_fn=denoised_fn,
636
+ model_kwargs=model_kwargs,
637
+ eta=eta,
638
+ )
639
+ yield out
640
+ img = out["sample"]
641
+
642
+ def _vb_terms_bpd(
643
+ self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None
644
+ ):
645
+ """
646
+ Get a term for the variational lower-bound.
647
+
648
+ The resulting units are bits (rather than nats, as one might expect).
649
+ This allows for comparison to other papers.
650
+
651
+ :return: a dict with the following keys:
652
+ - 'output': a shape [N] tensor of NLLs or KLs.
653
+ - 'pred_xstart': the x_0 predictions.
654
+ """
655
+ true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
656
+ x_start=x_start, x_t=x_t, t=t
657
+ )
658
+ out = self.p_mean_variance(
659
+ model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs
660
+ )
661
+ kl = normal_kl(
662
+ true_mean, true_log_variance_clipped, out["mean"], out["log_variance"]
663
+ )
664
+ kl = mean_flat(kl) / np.log(2.0)
665
+
666
+ decoder_nll = -discretized_gaussian_log_likelihood(
667
+ x_start, means=out["mean"], log_scales=0.5 * out["log_variance"]
668
+ )
669
+ assert decoder_nll.shape == x_start.shape
670
+ decoder_nll = mean_flat(decoder_nll) / np.log(2.0)
671
+
672
+ # At the first timestep return the decoder NLL,
673
+ # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
674
+ output = th.where((t == 0), decoder_nll, kl)
675
+ return {"output": output, "pred_xstart": out["pred_xstart"]}
676
+
677
+ def training_losses(self, model, x_start, t, model_kwargs=None, noise=None):
678
+ """
679
+ Compute training losses for a single timestep.
680
+
681
+ :param model: the model to evaluate loss on.
682
+ :param x_start: the [N x C x ...] tensor of inputs.
683
+ :param t: a batch of timestep indices.
684
+ :param model_kwargs: if not None, a dict of extra keyword arguments to
685
+ pass to the model. This can be used for conditioning.
686
+ :param noise: if specified, the specific Gaussian noise to try to remove.
687
+ :return: a dict with the key "loss" containing a tensor of shape [N].
688
+ Some mean or variance settings may also have other keys.
689
+ """
690
+ if model_kwargs is None:
691
+ model_kwargs = {}
692
+ if noise is None:
693
+ noise = th.randn_like(x_start)
694
+ x_t = self.q_sample(x_start, t, noise=noise)
695
+
696
+ terms = {}
697
+
698
+ if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL:
699
+ terms["loss"] = self._vb_terms_bpd(
700
+ model=model,
701
+ x_start=x_start,
702
+ x_t=x_t,
703
+ t=t,
704
+ clip_denoised=False,
705
+ model_kwargs=model_kwargs,
706
+ )["output"]
707
+ if self.loss_type == LossType.RESCALED_KL:
708
+ terms["loss"] *= self.num_timesteps
709
+ elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE:
710
+ model_output = model(x_t, self._scale_timesteps(t), **model_kwargs)
711
+
712
+ if self.model_var_type in [
713
+ ModelVarType.LEARNED,
714
+ ModelVarType.LEARNED_RANGE,
715
+ ]:
716
+ B, C = x_t.shape[:2]
717
+ assert model_output.shape == (B, C * 2, *x_t.shape[2:])
718
+ model_output, model_var_values = th.split(model_output, C, dim=1)
719
+ # Learn the variance using the variational bound, but don't let
720
+ # it affect our mean prediction.
721
+ frozen_out = th.cat([model_output.detach(), model_var_values], dim=1)
722
+ terms["vb"] = self._vb_terms_bpd(
723
+ model=lambda *args, r=frozen_out: r,
724
+ x_start=x_start,
725
+ x_t=x_t,
726
+ t=t,
727
+ clip_denoised=False,
728
+ )["output"]
729
+ if self.loss_type == LossType.RESCALED_MSE:
730
+ # Divide by 1000 for equivalence with initial implementation.
731
+ # Without a factor of 1/1000, the VB term hurts the MSE term.
732
+ terms["vb"] *= self.num_timesteps / 1000.0
733
+
734
+ target = {
735
+ ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance(
736
+ x_start=x_start, x_t=x_t, t=t
737
+ )[0],
738
+ ModelMeanType.START_X: x_start,
739
+ ModelMeanType.EPSILON: noise,
740
+ }[self.model_mean_type]
741
+ assert model_output.shape == target.shape == x_start.shape
742
+ terms["mse"] = mean_flat((target - model_output) ** 2)
743
+ if "vb" in terms:
744
+ terms["loss"] = terms["mse"] + terms["vb"]
745
+ else:
746
+ terms["loss"] = terms["mse"]
747
+ else:
748
+ raise NotImplementedError(self.loss_type)
749
+
750
+ return terms
751
+
752
+ def _prior_bpd(self, x_start):
753
+ """
754
+ Get the prior KL term for the variational lower-bound, measured in
755
+ bits-per-dim.
756
+
757
+ This term can't be optimized, as it only depends on the encoder.
758
+
759
+ :param x_start: the [N x C x ...] tensor of inputs.
760
+ :return: a batch of [N] KL values (in bits), one per batch element.
761
+ """
762
+ batch_size = x_start.shape[0]
763
+ t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
764
+ qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
765
+ kl_prior = normal_kl(
766
+ mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0
767
+ )
768
+ return mean_flat(kl_prior) / np.log(2.0)
769
+
770
+ def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None):
771
+ """
772
+ Compute the entire variational lower-bound, measured in bits-per-dim,
773
+ as well as other related quantities.
774
+
775
+ :param model: the model to evaluate loss on.
776
+ :param x_start: the [N x C x ...] tensor of inputs.
777
+ :param clip_denoised: if True, clip denoised samples.
778
+ :param model_kwargs: if not None, a dict of extra keyword arguments to
779
+ pass to the model. This can be used for conditioning.
780
+
781
+ :return: a dict containing the following keys:
782
+ - total_bpd: the total variational lower-bound, per batch element.
783
+ - prior_bpd: the prior term in the lower-bound.
784
+ - vb: an [N x T] tensor of terms in the lower-bound.
785
+ - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep.
786
+ - mse: an [N x T] tensor of epsilon MSEs for each timestep.
787
+ """
788
+ device = x_start.device
789
+ batch_size = x_start.shape[0]
790
+
791
+ vb = []
792
+ xstart_mse = []
793
+ mse = []
794
+ for t in list(range(self.num_timesteps))[::-1]:
795
+ t_batch = th.tensor([t] * batch_size, device=device)
796
+ noise = th.randn_like(x_start)
797
+ x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise)
798
+ # Calculate VLB term at the current timestep
799
+ with th.no_grad():
800
+ out = self._vb_terms_bpd(
801
+ model,
802
+ x_start=x_start,
803
+ x_t=x_t,
804
+ t=t_batch,
805
+ clip_denoised=clip_denoised,
806
+ model_kwargs=model_kwargs,
807
+ )
808
+ vb.append(out["output"])
809
+ xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2))
810
+ eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"])
811
+ mse.append(mean_flat((eps - noise) ** 2))
812
+
813
+ vb = th.stack(vb, dim=1)
814
+ xstart_mse = th.stack(xstart_mse, dim=1)
815
+ mse = th.stack(mse, dim=1)
816
+
817
+ prior_bpd = self._prior_bpd(x_start)
818
+ total_bpd = vb.sum(dim=1) + prior_bpd
819
+ return {
820
+ "total_bpd": total_bpd,
821
+ "prior_bpd": prior_bpd,
822
+ "vb": vb,
823
+ "xstart_mse": xstart_mse,
824
+ "mse": mse,
825
+ }
826
+
827
+
828
+ def _extract_into_tensor(arr, timesteps, broadcast_shape):
829
+ """
830
+ Extract values from a 1-D numpy array for a batch of indices.
831
+
832
+ :param arr: the 1-D numpy array.
833
+ :param timesteps: a tensor of indices into the array to extract.
834
+ :param broadcast_shape: a larger shape of K dimensions with the batch
835
+ dimension equal to the length of timesteps.
836
+ :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
837
+ """
838
+ res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
839
+ while len(res.shape) < len(broadcast_shape):
840
+ res = res[..., None]
841
+ return res.expand(broadcast_shape)
improved_diffusion/image_datasets.py ADDED
@@ -0,0 +1,106 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from PIL import Image
2
+ import blobfile as bf
3
+ from mpi4py import MPI
4
+ import numpy as np
5
+ from torch.utils.data import DataLoader, Dataset
6
+
7
+
8
+ def load_data(
9
+ *, data_dir, batch_size, image_size, class_cond=False, deterministic=False
10
+ ):
11
+ """
12
+ For a dataset, create a generator over (images, kwargs) pairs.
13
+
14
+ Each images is an NCHW float tensor, and the kwargs dict contains zero or
15
+ more keys, each of which map to a batched Tensor of their own.
16
+ The kwargs dict can be used for class labels, in which case the key is "y"
17
+ and the values are integer tensors of class labels.
18
+
19
+ :param data_dir: a dataset directory.
20
+ :param batch_size: the batch size of each returned pair.
21
+ :param image_size: the size to which images are resized.
22
+ :param class_cond: if True, include a "y" key in returned dicts for class
23
+ label. If classes are not available and this is true, an
24
+ exception will be raised.
25
+ :param deterministic: if True, yield results in a deterministic order.
26
+ """
27
+ if not data_dir:
28
+ raise ValueError("unspecified data directory")
29
+ all_files = _list_image_files_recursively(data_dir)
30
+ classes = None
31
+ if class_cond:
32
+ # Assume classes are the first part of the filename,
33
+ # before an underscore.
34
+ class_names = [bf.basename(path).split("_")[0] for path in all_files]
35
+ sorted_classes = {x: i for i, x in enumerate(sorted(set(class_names)))}
36
+ classes = [sorted_classes[x] for x in class_names]
37
+ dataset = ImageDataset(
38
+ image_size,
39
+ all_files,
40
+ classes=classes,
41
+ shard=MPI.COMM_WORLD.Get_rank(),
42
+ num_shards=MPI.COMM_WORLD.Get_size(),
43
+ )
44
+ if deterministic:
45
+ loader = DataLoader(
46
+ dataset, batch_size=batch_size, shuffle=False, num_workers=1, drop_last=True
47
+ )
48
+ else:
49
+ loader = DataLoader(
50
+ dataset, batch_size=batch_size, shuffle=True, num_workers=1, drop_last=True
51
+ )
52
+ while True:
53
+ yield from loader
54
+
55
+
56
+ def _list_image_files_recursively(data_dir):
57
+ results = []
58
+ for entry in sorted(bf.listdir(data_dir)):
59
+ full_path = bf.join(data_dir, entry)
60
+ ext = entry.split(".")[-1]
61
+ if "." in entry and ext.lower() in ["jpg", "jpeg", "png", "gif"]:
62
+ results.append(full_path)
63
+ elif bf.isdir(full_path):
64
+ results.extend(_list_image_files_recursively(full_path))
65
+ return results
66
+
67
+
68
+ class ImageDataset(Dataset):
69
+ def __init__(self, resolution, image_paths, classes=None, shard=0, num_shards=1):
70
+ super().__init__()
71
+ self.resolution = resolution
72
+ self.local_images = image_paths[shard:][::num_shards]
73
+ self.local_classes = None if classes is None else classes[shard:][::num_shards]
74
+
75
+ def __len__(self):
76
+ return len(self.local_images)
77
+
78
+ def __getitem__(self, idx):
79
+ path = self.local_images[idx]
80
+ with bf.BlobFile(path, "rb") as f:
81
+ pil_image = Image.open(f)
82
+ pil_image.load()
83
+
84
+ # We are not on a new enough PIL to support the `reducing_gap`
85
+ # argument, which uses BOX downsampling at powers of two first.
86
+ # Thus, we do it by hand to improve downsample quality.
87
+ while min(*pil_image.size) >= 2 * self.resolution:
88
+ pil_image = pil_image.resize(
89
+ tuple(x // 2 for x in pil_image.size), resample=Image.BOX
90
+ )
91
+
92
+ scale = self.resolution / min(*pil_image.size)
93
+ pil_image = pil_image.resize(
94
+ tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
95
+ )
96
+
97
+ arr = np.array(pil_image.convert("RGB"))
98
+ crop_y = (arr.shape[0] - self.resolution) // 2
99
+ crop_x = (arr.shape[1] - self.resolution) // 2
100
+ arr = arr[crop_y : crop_y + self.resolution, crop_x : crop_x + self.resolution]
101
+ arr = arr.astype(np.float32) / 127.5 - 1
102
+
103
+ out_dict = {}
104
+ if self.local_classes is not None:
105
+ out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64)
106
+ return np.transpose(arr, [2, 0, 1]), out_dict
improved_diffusion/logger.py ADDED
@@ -0,0 +1,495 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+
improved_diffusion/losses.py ADDED
@@ -0,0 +1,77 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
improved_diffusion/nn.py ADDED
@@ -0,0 +1,170 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+
10
+
11
+ # PyTorch 1.7 has SiLU, but we support PyTorch 1.5.
12
+ class SiLU(nn.Module):
13
+ def forward(self, x):
14
+ return x * th.sigmoid(x)
15
+
16
+
17
+ class GroupNorm32(nn.GroupNorm):
18
+ def forward(self, x):
19
+ return super().forward(x.float()).type(x.dtype)
20
+
21
+
22
+ def conv_nd(dims, *args, **kwargs):
23
+ """
24
+ Create a 1D, 2D, or 3D convolution module.
25
+ """
26
+ if dims == 1:
27
+ return nn.Conv1d(*args, **kwargs)
28
+ elif dims == 2:
29
+ return nn.Conv2d(*args, **kwargs)
30
+ elif dims == 3:
31
+ return nn.Conv3d(*args, **kwargs)
32
+ raise ValueError(f"unsupported dimensions: {dims}")
33
+
34
+
35
+ def linear(*args, **kwargs):
36
+ """
37
+ Create a linear module.
38
+ """
39
+ return nn.Linear(*args, **kwargs)
40
+
41
+
42
+ def avg_pool_nd(dims, *args, **kwargs):
43
+ """
44
+ Create a 1D, 2D, or 3D average pooling module.
45
+ """
46
+ if dims == 1:
47
+ return nn.AvgPool1d(*args, **kwargs)
48
+ elif dims == 2:
49
+ return nn.AvgPool2d(*args, **kwargs)
50
+ elif dims == 3:
51
+ return nn.AvgPool3d(*args, **kwargs)
52
+ raise ValueError(f"unsupported dimensions: {dims}")
53
+
54
+
55
+ def update_ema(target_params, source_params, rate=0.99):
56
+ """
57
+ Update target parameters to be closer to those of source parameters using
58
+ an exponential moving average.
59
+
60
+ :param target_params: the target parameter sequence.
61
+ :param source_params: the source parameter sequence.
62
+ :param rate: the EMA rate (closer to 1 means slower).
63
+ """
64
+ for targ, src in zip(target_params, source_params):
65
+ targ.detach().mul_(rate).add_(src, alpha=1 - rate)
66
+
67
+
68
+ def zero_module(module):
69
+ """
70
+ Zero out the parameters of a module and return it.
71
+ """
72
+ for p in module.parameters():
73
+ p.detach().zero_()
74
+ return module
75
+
76
+
77
+ def scale_module(module, scale):
78
+ """
79
+ Scale the parameters of a module and return it.
80
+ """
81
+ for p in module.parameters():
82
+ p.detach().mul_(scale)
83
+ return module
84
+
85
+
86
+ def mean_flat(tensor):
87
+ """
88
+ Take the mean over all non-batch dimensions.
89
+ """
90
+ return tensor.mean(dim=list(range(1, len(tensor.shape))))
91
+
92
+
93
+ def normalization(channels):
94
+ """
95
+ Make a standard normalization layer.
96
+
97
+ :param channels: number of input channels.
98
+ :return: an nn.Module for normalization.
99
+ """
100
+ return GroupNorm32(32, channels)
101
+
102
+
103
+ def timestep_embedding(timesteps, dim, max_period=10000):
104
+ """
105
+ Create sinusoidal timestep embeddings.
106
+
107
+ :param timesteps: a 1-D Tensor of N indices, one per batch element.
108
+ These may be fractional.
109
+ :param dim: the dimension of the output.
110
+ :param max_period: controls the minimum frequency of the embeddings.
111
+ :return: an [N x dim] Tensor of positional embeddings.
112
+ """
113
+ half = dim // 2
114
+ freqs = th.exp(
115
+ -math.log(max_period) * th.arange(start=0, end=half, dtype=th.float32) / half
116
+ ).to(device=timesteps.device)
117
+ args = timesteps[:, None].float() * freqs[None]
118
+ embedding = th.cat([th.cos(args), th.sin(args)], dim=-1)
119
+ if dim % 2:
120
+ embedding = th.cat([embedding, th.zeros_like(embedding[:, :1])], dim=-1)
121
+ return embedding
122
+
123
+
124
+ def checkpoint(func, inputs, params, flag):
125
+ """
126
+ Evaluate a function without caching intermediate activations, allowing for
127
+ reduced memory at the expense of extra compute in the backward pass.
128
+
129
+ :param func: the function to evaluate.
130
+ :param inputs: the argument sequence to pass to `func`.
131
+ :param params: a sequence of parameters `func` depends on but does not
132
+ explicitly take as arguments.
133
+ :param flag: if False, disable gradient checkpointing.
134
+ """
135
+ if flag:
136
+ args = tuple(inputs) + tuple(params)
137
+ return CheckpointFunction.apply(func, len(inputs), *args)
138
+ else:
139
+ return func(*inputs)
140
+
141
+
142
+ class CheckpointFunction(th.autograd.Function):
143
+ @staticmethod
144
+ def forward(ctx, run_function, length, *args):
145
+ ctx.run_function = run_function
146
+ ctx.input_tensors = list(args[:length])
147
+ ctx.input_params = list(args[length:])
148
+ with th.no_grad():
149
+ output_tensors = ctx.run_function(*ctx.input_tensors)
150
+ return output_tensors
151
+
152
+ @staticmethod
153
+ def backward(ctx, *output_grads):
154
+ ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
155
+ with th.enable_grad():
156
+ # Fixes a bug where the first op in run_function modifies the
157
+ # Tensor storage in place, which is not allowed for detach()'d
158
+ # Tensors.
159
+ shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
160
+ output_tensors = ctx.run_function(*shallow_copies)
161
+ input_grads = th.autograd.grad(
162
+ output_tensors,
163
+ ctx.input_tensors + ctx.input_params,
164
+ output_grads,
165
+ allow_unused=True,
166
+ )
167
+ del ctx.input_tensors
168
+ del ctx.input_params
169
+ del output_tensors
170
+ return (None, None) + input_grads
improved_diffusion/resample.py ADDED
@@ -0,0 +1,154 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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()
improved_diffusion/respace.py ADDED
@@ -0,0 +1,122 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 _wrap_model(self, model):
99
+ if isinstance(model, _WrappedModel):
100
+ return model
101
+ return _WrappedModel(
102
+ model, self.timestep_map, self.rescale_timesteps, self.original_num_steps
103
+ )
104
+
105
+ def _scale_timesteps(self, t):
106
+ # Scaling is done by the wrapped model.
107
+ return t
108
+
109
+
110
+ class _WrappedModel:
111
+ def __init__(self, model, timestep_map, rescale_timesteps, original_num_steps):
112
+ self.model = model
113
+ self.timestep_map = timestep_map
114
+ self.rescale_timesteps = rescale_timesteps
115
+ self.original_num_steps = original_num_steps
116
+
117
+ def __call__(self, x, ts, **kwargs):
118
+ map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)
119
+ new_ts = map_tensor[ts]
120
+ if self.rescale_timesteps:
121
+ new_ts = new_ts.float() * (1000.0 / self.original_num_steps)
122
+ return self.model(x, new_ts, **kwargs)
improved_diffusion/script_util.py ADDED
@@ -0,0 +1,296 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 SuperResModel, UNetModel
7
+
8
+ NUM_CLASSES = 1000
9
+
10
+
11
+ def model_and_diffusion_defaults():
12
+ """
13
+ Defaults for image training.
14
+ """
15
+ return dict(
16
+ image_size=64,
17
+ num_channels=128,
18
+ num_res_blocks=2,
19
+ num_heads=4,
20
+ num_heads_upsample=-1,
21
+ attention_resolutions="16,8",
22
+ dropout=0.0,
23
+ learn_sigma=False,
24
+ sigma_small=False,
25
+ class_cond=False,
26
+ diffusion_steps=1000,
27
+ noise_schedule="linear",
28
+ timestep_respacing="",
29
+ use_kl=False,
30
+ predict_xstart=False,
31
+ rescale_timesteps=True,
32
+ rescale_learned_sigmas=True,
33
+ use_checkpoint=False,
34
+ use_scale_shift_norm=True,
35
+ )
36
+
37
+
38
+ def create_model_and_diffusion(
39
+ image_size,
40
+ class_cond,
41
+ learn_sigma,
42
+ sigma_small,
43
+ num_channels,
44
+ num_res_blocks,
45
+ num_heads,
46
+ num_heads_upsample,
47
+ attention_resolutions,
48
+ dropout,
49
+ diffusion_steps,
50
+ noise_schedule,
51
+ timestep_respacing,
52
+ use_kl,
53
+ predict_xstart,
54
+ rescale_timesteps,
55
+ rescale_learned_sigmas,
56
+ use_checkpoint,
57
+ use_scale_shift_norm,
58
+ ):
59
+ model = create_model(
60
+ image_size,
61
+ num_channels,
62
+ num_res_blocks,
63
+ learn_sigma=learn_sigma,
64
+ class_cond=class_cond,
65
+ use_checkpoint=use_checkpoint,
66
+ attention_resolutions=attention_resolutions,
67
+ num_heads=num_heads,
68
+ num_heads_upsample=num_heads_upsample,
69
+ use_scale_shift_norm=use_scale_shift_norm,
70
+ dropout=dropout,
71
+ )
72
+ diffusion = create_gaussian_diffusion(
73
+ steps=diffusion_steps,
74
+ learn_sigma=learn_sigma,
75
+ sigma_small=sigma_small,
76
+ noise_schedule=noise_schedule,
77
+ use_kl=use_kl,
78
+ predict_xstart=predict_xstart,
79
+ rescale_timesteps=rescale_timesteps,
80
+ rescale_learned_sigmas=rescale_learned_sigmas,
81
+ timestep_respacing=timestep_respacing,
82
+ )
83
+ return model, diffusion
84
+
85
+
86
+ def create_model(
87
+ image_size,
88
+ num_channels,
89
+ num_res_blocks,
90
+ learn_sigma,
91
+ class_cond,
92
+ use_checkpoint,
93
+ attention_resolutions,
94
+ num_heads,
95
+ num_heads_upsample,
96
+ use_scale_shift_norm,
97
+ dropout,
98
+ ):
99
+ if image_size == 256:
100
+ channel_mult = (1, 1, 2, 2, 4, 4)
101
+ elif image_size == 64:
102
+ channel_mult = (1, 2, 3, 4)
103
+ elif image_size == 32:
104
+ channel_mult = (1, 2, 2, 2)
105
+ else:
106
+ raise ValueError(f"unsupported image size: {image_size}")
107
+
108
+ attention_ds = []
109
+ for res in attention_resolutions.split(","):
110
+ attention_ds.append(image_size // int(res))
111
+
112
+ return UNetModel(
113
+ in_channels=3,
114
+ model_channels=num_channels,
115
+ out_channels=(3 if not learn_sigma else 6),
116
+ num_res_blocks=num_res_blocks,
117
+ attention_resolutions=tuple(attention_ds),
118
+ dropout=dropout,
119
+ channel_mult=channel_mult,
120
+ num_classes=(NUM_CLASSES if class_cond else None),
121
+ use_checkpoint=use_checkpoint,
122
+ num_heads=num_heads,
123
+ num_heads_upsample=num_heads_upsample,
124
+ use_scale_shift_norm=use_scale_shift_norm,
125
+ )
126
+
127
+
128
+ def sr_model_and_diffusion_defaults():
129
+ res = model_and_diffusion_defaults()
130
+ res["large_size"] = 256
131
+ res["small_size"] = 64
132
+ arg_names = inspect.getfullargspec(sr_create_model_and_diffusion)[0]
133
+ for k in res.copy().keys():
134
+ if k not in arg_names:
135
+ del res[k]
136
+ return res
137
+
138
+
139
+ def sr_create_model_and_diffusion(
140
+ large_size,
141
+ small_size,
142
+ class_cond,
143
+ learn_sigma,
144
+ num_channels,
145
+ num_res_blocks,
146
+ num_heads,
147
+ num_heads_upsample,
148
+ attention_resolutions,
149
+ dropout,
150
+ diffusion_steps,
151
+ noise_schedule,
152
+ timestep_respacing,
153
+ use_kl,
154
+ predict_xstart,
155
+ rescale_timesteps,
156
+ rescale_learned_sigmas,
157
+ use_checkpoint,
158
+ use_scale_shift_norm,
159
+ ):
160
+ model = sr_create_model(
161
+ large_size,
162
+ small_size,
163
+ num_channels,
164
+ num_res_blocks,
165
+ learn_sigma=learn_sigma,
166
+ class_cond=class_cond,
167
+ use_checkpoint=use_checkpoint,
168
+ attention_resolutions=attention_resolutions,
169
+ num_heads=num_heads,
170
+ num_heads_upsample=num_heads_upsample,
171
+ use_scale_shift_norm=use_scale_shift_norm,
172
+ dropout=dropout,
173
+ )
174
+ diffusion = create_gaussian_diffusion(
175
+ steps=diffusion_steps,
176
+ learn_sigma=learn_sigma,
177
+ noise_schedule=noise_schedule,
178
+ use_kl=use_kl,
179
+ predict_xstart=predict_xstart,
180
+ rescale_timesteps=rescale_timesteps,
181
+ rescale_learned_sigmas=rescale_learned_sigmas,
182
+ timestep_respacing=timestep_respacing,
183
+ )
184
+ return model, diffusion
185
+
186
+
187
+ def sr_create_model(
188
+ large_size,
189
+ small_size,
190
+ num_channels,
191
+ num_res_blocks,
192
+ learn_sigma,
193
+ class_cond,
194
+ use_checkpoint,
195
+ attention_resolutions,
196
+ num_heads,
197
+ num_heads_upsample,
198
+ use_scale_shift_norm,
199
+ dropout,
200
+ ):
201
+ _ = small_size # hack to prevent unused variable
202
+
203
+ if large_size == 256:
204
+ channel_mult = (1, 1, 2, 2, 4, 4)
205
+ elif large_size == 64:
206
+ channel_mult = (1, 2, 3, 4)
207
+ else:
208
+ raise ValueError(f"unsupported large size: {large_size}")
209
+
210
+ attention_ds = []
211
+ for res in attention_resolutions.split(","):
212
+ attention_ds.append(large_size // int(res))
213
+
214
+ return SuperResModel(
215
+ in_channels=3,
216
+ model_channels=num_channels,
217
+ out_channels=(3 if not learn_sigma else 6),
218
+ num_res_blocks=num_res_blocks,
219
+ attention_resolutions=tuple(attention_ds),
220
+ dropout=dropout,
221
+ channel_mult=channel_mult,
222
+ num_classes=(NUM_CLASSES if class_cond else None),
223
+ use_checkpoint=use_checkpoint,
224
+ num_heads=num_heads,
225
+ num_heads_upsample=num_heads_upsample,
226
+ use_scale_shift_norm=use_scale_shift_norm,
227
+ )
228
+
229
+
230
+ def create_gaussian_diffusion(
231
+ *,
232
+ steps=1000,
233
+ learn_sigma=False,
234
+ sigma_small=False,
235
+ noise_schedule="linear",
236
+ use_kl=False,
237
+ predict_xstart=False,
238
+ rescale_timesteps=False,
239
+ rescale_learned_sigmas=False,
240
+ timestep_respacing="",
241
+ ):
242
+ betas = gd.get_named_beta_schedule(noise_schedule, steps)
243
+ if use_kl:
244
+ loss_type = gd.LossType.RESCALED_KL
245
+ elif rescale_learned_sigmas:
246
+ loss_type = gd.LossType.RESCALED_MSE
247
+ else:
248
+ loss_type = gd.LossType.MSE
249
+ if not timestep_respacing:
250
+ timestep_respacing = [steps]
251
+ return SpacedDiffusion(
252
+ use_timesteps=space_timesteps(steps, timestep_respacing),
253
+ betas=betas,
254
+ model_mean_type=(
255
+ gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X
256
+ ),
257
+ model_var_type=(
258
+ (
259
+ gd.ModelVarType.FIXED_LARGE
260
+ if not sigma_small
261
+ else gd.ModelVarType.FIXED_SMALL
262
+ )
263
+ if not learn_sigma
264
+ else gd.ModelVarType.LEARNED_RANGE
265
+ ),
266
+ loss_type=loss_type,
267
+ rescale_timesteps=rescale_timesteps,
268
+ )
269
+
270
+
271
+ def add_dict_to_argparser(parser, default_dict):
272
+ for k, v in default_dict.items():
273
+ v_type = type(v)
274
+ if v is None:
275
+ v_type = str
276
+ elif isinstance(v, bool):
277
+ v_type = str2bool
278
+ parser.add_argument(f"--{k}", default=v, type=v_type)
279
+
280
+
281
+ def args_to_dict(args, keys):
282
+ return {k: getattr(args, k) for k in keys}
283
+
284
+
285
+ def str2bool(v):
286
+ """
287
+ https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse
288
+ """
289
+ if isinstance(v, bool):
290
+ return v
291
+ if v.lower() in ("yes", "true", "t", "y", "1"):
292
+ return True
293
+ elif v.lower() in ("no", "false", "f", "n", "0"):
294
+ return False
295
+ else:
296
+ raise argparse.ArgumentTypeError("boolean value expected")
improved_diffusion/train_util.py ADDED
@@ -0,0 +1,356 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import copy
2
+ import functools
3
+ import os
4
+
5
+ import blobfile as bf
6
+ import numpy as np
7
+ import torch as th
8
+ import torch.distributed as dist
9
+ from torch.nn.parallel.distributed import DistributedDataParallel as DDP
10
+ from torch.optim import AdamW
11
+
12
+ from . import dist_util, logger
13
+ from .fp16_util import (
14
+ make_master_params,
15
+ master_params_to_model_params,
16
+ model_grads_to_master_grads,
17
+ unflatten_master_params,
18
+ zero_grad,
19
+ )
20
+ from .nn import update_ema
21
+ from .resample import LossAwareSampler, UniformSampler
22
+
23
+ # For ImageNet experiments, this was a good default value.
24
+ # We found that the lg_loss_scale quickly climbed to
25
+ # 20-21 within the first ~1K steps of training.
26
+ INITIAL_LOG_LOSS_SCALE = 20.0
27
+
28
+
29
+ class TrainLoop:
30
+ def __init__(
31
+ self,
32
+ *,
33
+ model,
34
+ diffusion,
35
+ data,
36
+ batch_size,
37
+ microbatch,
38
+ lr,
39
+ ema_rate,
40
+ log_interval,
41
+ save_interval,
42
+ resume_checkpoint,
43
+ use_fp16=False,
44
+ fp16_scale_growth=1e-3,
45
+ schedule_sampler=None,
46
+ weight_decay=0.0,
47
+ lr_anneal_steps=0,
48
+ ):
49
+ self.model = model
50
+ self.diffusion = diffusion
51
+ self.data = data
52
+ self.batch_size = batch_size
53
+ self.microbatch = microbatch if microbatch > 0 else batch_size
54
+ self.lr = lr
55
+ self.ema_rate = (
56
+ [ema_rate]
57
+ if isinstance(ema_rate, float)
58
+ else [float(x) for x in ema_rate.split(",")]
59
+ )
60
+ self.log_interval = log_interval
61
+ self.save_interval = save_interval
62
+ self.resume_checkpoint = resume_checkpoint
63
+ self.use_fp16 = use_fp16
64
+ self.fp16_scale_growth = fp16_scale_growth
65
+ self.schedule_sampler = schedule_sampler or UniformSampler(diffusion)
66
+ self.weight_decay = weight_decay
67
+ self.lr_anneal_steps = lr_anneal_steps
68
+
69
+ self.step = 0
70
+ self.resume_step = 0
71
+ self.global_batch = self.batch_size * dist.get_world_size()
72
+
73
+ self.model_params = list(self.model.parameters())
74
+ self.master_params = self.model_params
75
+ self.lg_loss_scale = INITIAL_LOG_LOSS_SCALE
76
+ self.sync_cuda = th.cuda.is_available()
77
+
78
+ self._load_and_sync_parameters()
79
+ if self.use_fp16:
80
+ self._setup_fp16()
81
+
82
+ self.opt = AdamW(self.master_params, lr=self.lr, weight_decay=self.weight_decay)
83
+ if self.resume_step:
84
+ self._load_optimizer_state()
85
+ # Model was resumed, either due to a restart or a checkpoint
86
+ # being specified at the command line.
87
+ self.ema_params = [
88
+ self._load_ema_parameters(rate) for rate in self.ema_rate
89
+ ]
90
+ else:
91
+ self.ema_params = [
92
+ copy.deepcopy(self.master_params) for _ in range(len(self.ema_rate))
93
+ ]
94
+
95
+ if th.cuda.is_available():
96
+ self.use_ddp = True
97
+ self.ddp_model = DDP(
98
+ self.model,
99
+ device_ids=[dist_util.dev()],
100
+ output_device=dist_util.dev(),
101
+ broadcast_buffers=False,
102
+ bucket_cap_mb=128,
103
+ find_unused_parameters=False,
104
+ )
105
+ else:
106
+ if dist.get_world_size() > 1:
107
+ logger.warn(
108
+ "Distributed training requires CUDA. "
109
+ "Gradients will not be synchronized properly!"
110
+ )
111
+ self.use_ddp = False
112
+ self.ddp_model = self.model
113
+
114
+ def _load_and_sync_parameters(self):
115
+ resume_checkpoint = find_resume_checkpoint() or self.resume_checkpoint
116
+
117
+ if resume_checkpoint:
118
+ self.resume_step = parse_resume_step_from_filename(resume_checkpoint)
119
+ if dist.get_rank() == 0:
120
+ logger.log(f"loading model from checkpoint: {resume_checkpoint}...")
121
+ self.model.load_state_dict(
122
+ dist_util.load_state_dict(
123
+ resume_checkpoint, map_location=dist_util.dev()
124
+ )
125
+ )
126
+
127
+ dist_util.sync_params(self.model.parameters())
128
+
129
+ def _load_ema_parameters(self, rate):
130
+ ema_params = copy.deepcopy(self.master_params)
131
+
132
+ main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint
133
+ ema_checkpoint = find_ema_checkpoint(main_checkpoint, self.resume_step, rate)
134
+ if ema_checkpoint:
135
+ if dist.get_rank() == 0:
136
+ logger.log(f"loading EMA from checkpoint: {ema_checkpoint}...")
137
+ state_dict = dist_util.load_state_dict(
138
+ ema_checkpoint, map_location=dist_util.dev()
139
+ )
140
+ ema_params = self._state_dict_to_master_params(state_dict)
141
+
142
+ dist_util.sync_params(ema_params)
143
+ return ema_params
144
+
145
+ def _load_optimizer_state(self):
146
+ main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint
147
+ opt_checkpoint = bf.join(
148
+ bf.dirname(main_checkpoint), f"opt{self.resume_step:06}.pt"
149
+ )
150
+ if bf.exists(opt_checkpoint):
151
+ logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}")
152
+ state_dict = dist_util.load_state_dict(
153
+ opt_checkpoint, map_location=dist_util.dev()
154
+ )
155
+ self.opt.load_state_dict(state_dict)
156
+
157
+ def _setup_fp16(self):
158
+ self.master_params = make_master_params(self.model_params)
159
+ self.model.convert_to_fp16()
160
+
161
+ def run_loop(self):
162
+ while (
163
+ not self.lr_anneal_steps
164
+ or self.step + self.resume_step < self.lr_anneal_steps
165
+ ):
166
+ batch, cond = next(self.data)
167
+ self.run_step(batch, cond)
168
+ if self.step % self.log_interval == 0:
169
+ logger.dumpkvs()
170
+ if self.step % self.save_interval == 0:
171
+ self.save()
172
+ # Run for a finite amount of time in integration tests.
173
+ if os.environ.get("DIFFUSION_TRAINING_TEST", "") and self.step > 0:
174
+ return
175
+ self.step += 1
176
+ # Save the last checkpoint if it wasn't already saved.
177
+ if (self.step - 1) % self.save_interval != 0:
178
+ self.save()
179
+
180
+ def run_step(self, batch, cond):
181
+ self.forward_backward(batch, cond)
182
+ if self.use_fp16:
183
+ self.optimize_fp16()
184
+ else:
185
+ self.optimize_normal()
186
+ self.log_step()
187
+
188
+ def forward_backward(self, batch, cond):
189
+ zero_grad(self.model_params)
190
+ for i in range(0, batch.shape[0], self.microbatch):
191
+ micro = batch[i : i + self.microbatch].to(dist_util.dev())
192
+ micro_cond = {
193
+ k: v[i : i + self.microbatch].to(dist_util.dev())
194
+ for k, v in cond.items()
195
+ }
196
+ last_batch = (i + self.microbatch) >= batch.shape[0]
197
+ t, weights = self.schedule_sampler.sample(micro.shape[0], dist_util.dev())
198
+
199
+ compute_losses = functools.partial(
200
+ self.diffusion.training_losses,
201
+ self.ddp_model,
202
+ micro,
203
+ t,
204
+ model_kwargs=micro_cond,
205
+ )
206
+
207
+ if last_batch or not self.use_ddp:
208
+ losses = compute_losses()
209
+ else:
210
+ with self.ddp_model.no_sync():
211
+ losses = compute_losses()
212
+
213
+ if isinstance(self.schedule_sampler, LossAwareSampler):
214
+ self.schedule_sampler.update_with_local_losses(
215
+ t, losses["loss"].detach()
216
+ )
217
+
218
+ loss = (losses["loss"] * weights).mean()
219
+ log_loss_dict(
220
+ self.diffusion, t, {k: v * weights for k, v in losses.items()}
221
+ )
222
+ if self.use_fp16:
223
+ loss_scale = 2 ** self.lg_loss_scale
224
+ (loss * loss_scale).backward()
225
+ else:
226
+ loss.backward()
227
+
228
+ def optimize_fp16(self):
229
+ if any(not th.isfinite(p.grad).all() for p in self.model_params):
230
+ self.lg_loss_scale -= 1
231
+ logger.log(f"Found NaN, decreased lg_loss_scale to {self.lg_loss_scale}")
232
+ return
233
+
234
+ model_grads_to_master_grads(self.model_params, self.master_params)
235
+ self.master_params[0].grad.mul_(1.0 / (2 ** self.lg_loss_scale))
236
+ self._log_grad_norm()
237
+ self._anneal_lr()
238
+ self.opt.step()
239
+ for rate, params in zip(self.ema_rate, self.ema_params):
240
+ update_ema(params, self.master_params, rate=rate)
241
+ master_params_to_model_params(self.model_params, self.master_params)
242
+ self.lg_loss_scale += self.fp16_scale_growth
243
+
244
+ def optimize_normal(self):
245
+ self._log_grad_norm()
246
+ self._anneal_lr()
247
+ self.opt.step()
248
+ for rate, params in zip(self.ema_rate, self.ema_params):
249
+ update_ema(params, self.master_params, rate=rate)
250
+
251
+ def _log_grad_norm(self):
252
+ sqsum = 0.0
253
+ for p in self.master_params:
254
+ sqsum += (p.grad ** 2).sum().item()
255
+ logger.logkv_mean("grad_norm", np.sqrt(sqsum))
256
+
257
+ def _anneal_lr(self):
258
+ if not self.lr_anneal_steps:
259
+ return
260
+ frac_done = (self.step + self.resume_step) / self.lr_anneal_steps
261
+ lr = self.lr * (1 - frac_done)
262
+ for param_group in self.opt.param_groups:
263
+ param_group["lr"] = lr
264
+
265
+ def log_step(self):
266
+ logger.logkv("step", self.step + self.resume_step)
267
+ logger.logkv("samples", (self.step + self.resume_step + 1) * self.global_batch)
268
+ if self.use_fp16:
269
+ logger.logkv("lg_loss_scale", self.lg_loss_scale)
270
+
271
+ def save(self):
272
+ def save_checkpoint(rate, params):
273
+ state_dict = self._master_params_to_state_dict(params)
274
+ if dist.get_rank() == 0:
275
+ logger.log(f"saving model {rate}...")
276
+ if not rate:
277
+ filename = f"model{(self.step+self.resume_step):06d}.pt"
278
+ else:
279
+ filename = f"ema_{rate}_{(self.step+self.resume_step):06d}.pt"
280
+ with bf.BlobFile(bf.join(get_blob_logdir(), filename), "wb") as f:
281
+ th.save(state_dict, f)
282
+
283
+ save_checkpoint(0, self.master_params)
284
+ for rate, params in zip(self.ema_rate, self.ema_params):
285
+ save_checkpoint(rate, params)
286
+
287
+ if dist.get_rank() == 0:
288
+ with bf.BlobFile(
289
+ bf.join(get_blob_logdir(), f"opt{(self.step+self.resume_step):06d}.pt"),
290
+ "wb",
291
+ ) as f:
292
+ th.save(self.opt.state_dict(), f)
293
+
294
+ dist.barrier()
295
+
296
+ def _master_params_to_state_dict(self, master_params):
297
+ if self.use_fp16:
298
+ master_params = unflatten_master_params(
299
+ self.model.parameters(), master_params
300
+ )
301
+ state_dict = self.model.state_dict()
302
+ for i, (name, _value) in enumerate(self.model.named_parameters()):
303
+ assert name in state_dict
304
+ state_dict[name] = master_params[i]
305
+ return state_dict
306
+
307
+ def _state_dict_to_master_params(self, state_dict):
308
+ params = [state_dict[name] for name, _ in self.model.named_parameters()]
309
+ if self.use_fp16:
310
+ return make_master_params(params)
311
+ else:
312
+ return params
313
+
314
+
315
+ def parse_resume_step_from_filename(filename):
316
+ """
317
+ Parse filenames of the form path/to/modelNNNNNN.pt, where NNNNNN is the
318
+ checkpoint's number of steps.
319
+ """
320
+ split = filename.split("model")
321
+ if len(split) < 2:
322
+ return 0
323
+ split1 = split[-1].split(".")[0]
324
+ try:
325
+ return int(split1)
326
+ except ValueError:
327
+ return 0
328
+
329
+
330
+ def get_blob_logdir():
331
+ return os.environ.get("DIFFUSION_BLOB_LOGDIR", logger.get_dir())
332
+
333
+
334
+ def find_resume_checkpoint():
335
+ # On your infrastructure, you may want to override this to automatically
336
+ # discover the latest checkpoint on your blob storage, etc.
337
+ return None
338
+
339
+
340
+ def find_ema_checkpoint(main_checkpoint, step, rate):
341
+ if main_checkpoint is None:
342
+ return None
343
+ filename = f"ema_{rate}_{(step):06d}.pt"
344
+ path = bf.join(bf.dirname(main_checkpoint), filename)
345
+ if bf.exists(path):
346
+ return path
347
+ return None
348
+
349
+
350
+ def log_loss_dict(diffusion, ts, losses):
351
+ for key, values in losses.items():
352
+ logger.logkv_mean(key, values.mean().item())
353
+ # Log the quantiles (four quartiles, in particular).
354
+ for sub_t, sub_loss in zip(ts.cpu().numpy(), values.detach().cpu().numpy()):
355
+ quartile = int(4 * sub_t / diffusion.num_timesteps)
356
+ logger.logkv_mean(f"{key}_q{quartile}", sub_loss)
improved_diffusion/unet.py ADDED
@@ -0,0 +1,547 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from abc import abstractmethod
2
+
3
+ import math
4
+
5
+ import numpy as np
6
+ import torch as th
7
+ import torch.nn as nn
8
+ import torch.nn.functional as F
9
+
10
+ from .fp16_util import convert_module_to_f16, convert_module_to_f32
11
+ from .nn import (
12
+ SiLU,
13
+ conv_nd,
14
+ linear,
15
+ avg_pool_nd,
16
+ zero_module,
17
+ normalization,
18
+ timestep_embedding,
19
+ checkpoint,
20
+ )
21
+
22
+
23
+ class TimestepBlock(nn.Module):
24
+ """
25
+ Any module where forward() takes timestep embeddings as a second argument.
26
+ """
27
+
28
+ @abstractmethod
29
+ def forward(self, x, emb):
30
+ """
31
+ Apply the module to `x` given `emb` timestep embeddings.
32
+ """
33
+
34
+
35
+ class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
36
+ """
37
+ A sequential module that passes timestep embeddings to the children that
38
+ support it as an extra input.
39
+ """
40
+
41
+ def forward(self, x, emb):
42
+ for layer in self:
43
+ if isinstance(layer, TimestepBlock):
44
+ x = layer(x, emb)
45
+ else:
46
+ x = layer(x)
47
+ return x
48
+
49
+
50
+ class Upsample(nn.Module):
51
+ """
52
+ An upsampling layer with an optional convolution.
53
+
54
+ :param channels: channels in the inputs and outputs.
55
+ :param use_conv: a bool determining if a convolution is applied.
56
+ :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
57
+ upsampling occurs in the inner-two dimensions.
58
+ """
59
+
60
+ def __init__(self, channels, use_conv, dims=2):
61
+ super().__init__()
62
+ self.channels = channels
63
+ self.use_conv = use_conv
64
+ self.dims = dims
65
+ if use_conv:
66
+ self.conv = conv_nd(dims, channels, channels, 3, padding=1)
67
+
68
+ def forward(self, x):
69
+ assert x.shape[1] == self.channels
70
+ if self.dims == 3:
71
+ x = F.interpolate(
72
+ x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
73
+ )
74
+ else:
75
+ x = F.interpolate(x, scale_factor=2, mode="nearest")
76
+ if self.use_conv:
77
+ x = self.conv(x)
78
+ return x
79
+
80
+
81
+ class Downsample(nn.Module):
82
+ """
83
+ A downsampling layer with an optional convolution.
84
+
85
+ :param channels: channels in the inputs and outputs.
86
+ :param use_conv: a bool determining if a convolution is applied.
87
+ :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
88
+ downsampling occurs in the inner-two dimensions.
89
+ """
90
+
91
+ def __init__(self, channels, use_conv, dims=2):
92
+ super().__init__()
93
+ self.channels = channels
94
+ self.use_conv = use_conv
95
+ self.dims = dims
96
+ stride = 2 if dims != 3 else (1, 2, 2)
97
+ if use_conv:
98
+ self.op = conv_nd(dims, channels, channels, 3, stride=stride, padding=1)
99
+ else:
100
+ self.op = avg_pool_nd(stride)
101
+
102
+ def forward(self, x):
103
+ assert x.shape[1] == self.channels
104
+ return self.op(x)
105
+
106
+
107
+ class ResBlock(TimestepBlock):
108
+ """
109
+ A residual block that can optionally change the number of channels.
110
+
111
+ :param channels: the number of input channels.
112
+ :param emb_channels: the number of timestep embedding channels.
113
+ :param dropout: the rate of dropout.
114
+ :param out_channels: if specified, the number of out channels.
115
+ :param use_conv: if True and out_channels is specified, use a spatial
116
+ convolution instead of a smaller 1x1 convolution to change the
117
+ channels in the skip connection.
118
+ :param dims: determines if the signal is 1D, 2D, or 3D.
119
+ :param use_checkpoint: if True, use gradient checkpointing on this module.
120
+ """
121
+
122
+ def __init__(
123
+ self,
124
+ channels,
125
+ emb_channels,
126
+ dropout,
127
+ out_channels=None,
128
+ use_conv=False,
129
+ use_scale_shift_norm=False,
130
+ dims=2,
131
+ use_checkpoint=False,
132
+ ):
133
+ super().__init__()
134
+ self.channels = channels
135
+ self.emb_channels = emb_channels
136
+ self.dropout = dropout
137
+ self.out_channels = out_channels or channels
138
+ self.use_conv = use_conv
139
+ self.use_checkpoint = use_checkpoint
140
+ self.use_scale_shift_norm = use_scale_shift_norm
141
+
142
+ self.in_layers = nn.Sequential(
143
+ normalization(channels),
144
+ SiLU(),
145
+ conv_nd(dims, channels, self.out_channels, 3, padding=1),
146
+ )
147
+ self.emb_layers = nn.Sequential(
148
+ SiLU(),
149
+ linear(
150
+ emb_channels,
151
+ 2 * self.out_channels if use_scale_shift_norm else self.out_channels,
152
+ ),
153
+ )
154
+ self.out_layers = nn.Sequential(
155
+ normalization(self.out_channels),
156
+ SiLU(),
157
+ nn.Dropout(p=dropout),
158
+ zero_module(
159
+ conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)
160
+ ),
161
+ )
162
+
163
+ if self.out_channels == channels:
164
+ self.skip_connection = nn.Identity()
165
+ elif use_conv:
166
+ self.skip_connection = conv_nd(
167
+ dims, channels, self.out_channels, 3, padding=1
168
+ )
169
+ else:
170
+ self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
171
+
172
+ def forward(self, x, emb):
173
+ """
174
+ Apply the block to a Tensor, conditioned on a timestep embedding.
175
+
176
+ :param x: an [N x C x ...] Tensor of features.
177
+ :param emb: an [N x emb_channels] Tensor of timestep embeddings.
178
+ :return: an [N x C x ...] Tensor of outputs.
179
+ """
180
+ return checkpoint(
181
+ self._forward, (x, emb), self.parameters(), self.use_checkpoint
182
+ )
183
+
184
+ def _forward(self, x, emb):
185
+ h = self.in_layers(x)
186
+ emb_out = self.emb_layers(emb).type(h.dtype)
187
+ while len(emb_out.shape) < len(h.shape):
188
+ emb_out = emb_out[..., None]
189
+ if self.use_scale_shift_norm:
190
+ out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
191
+ scale, shift = th.chunk(emb_out, 2, dim=1)
192
+ h = out_norm(h) * (1 + scale) + shift
193
+ h = out_rest(h)
194
+ else:
195
+ h = h + emb_out
196
+ h = self.out_layers(h)
197
+ return self.skip_connection(x) + h
198
+
199
+
200
+ class AttentionBlock(nn.Module):
201
+ """
202
+ An attention block that allows spatial positions to attend to each other.
203
+
204
+ Originally ported from here, but adapted to the N-d case.
205
+ https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
206
+ """
207
+
208
+ def __init__(self, channels, num_heads=1, use_checkpoint=False):
209
+ super().__init__()
210
+ self.channels = channels
211
+ self.num_heads = num_heads
212
+ self.use_checkpoint = use_checkpoint
213
+
214
+ self.norm = normalization(channels)
215
+ self.qkv = conv_nd(1, channels, channels * 3, 1)
216
+ self.attention = QKVAttention()
217
+ self.proj_out = zero_module(conv_nd(1, channels, channels, 1))
218
+
219
+ def forward(self, x):
220
+ return checkpoint(self._forward, (x,), self.parameters(), self.use_checkpoint)
221
+
222
+ def _forward(self, x):
223
+ b, c, *spatial = x.shape
224
+ x = x.reshape(b, c, -1)
225
+ qkv = self.qkv(self.norm(x))
226
+ qkv = qkv.reshape(b * self.num_heads, -1, qkv.shape[2])
227
+ h = self.attention(qkv)
228
+ h = h.reshape(b, -1, h.shape[-1])
229
+ h = self.proj_out(h)
230
+ return (x + h).reshape(b, c, *spatial)
231
+
232
+
233
+ class QKVAttention(nn.Module):
234
+ """
235
+ A module which performs QKV attention.
236
+ """
237
+
238
+ def forward(self, qkv):
239
+ """
240
+ Apply QKV attention.
241
+
242
+ :param qkv: an [N x (C * 3) x T] tensor of Qs, Ks, and Vs.
243
+ :return: an [N x C x T] tensor after attention.
244
+ """
245
+ ch = qkv.shape[1] // 3
246
+ q, k, v = th.split(qkv, ch, dim=1)
247
+ scale = 1 / math.sqrt(math.sqrt(ch))
248
+ weight = th.einsum(
249
+ "bct,bcs->bts", q * scale, k * scale
250
+ ) # More stable with f16 than dividing afterwards
251
+ weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
252
+ return th.einsum("bts,bcs->bct", weight, v)
253
+
254
+ @staticmethod
255
+ def count_flops(model, _x, y):
256
+ """
257
+ A counter for the `thop` package to count the operations in an
258
+ attention operation.
259
+
260
+ Meant to be used like:
261
+
262
+ macs, params = thop.profile(
263
+ model,
264
+ inputs=(inputs, timestamps),
265
+ custom_ops={QKVAttention: QKVAttention.count_flops},
266
+ )
267
+
268
+ """
269
+ b, c, *spatial = y[0].shape
270
+ num_spatial = int(np.prod(spatial))
271
+ # We perform two matmuls with the same number of ops.
272
+ # The first computes the weight matrix, the second computes
273
+ # the combination of the value vectors.
274
+ matmul_ops = 2 * b * (num_spatial ** 2) * c
275
+ model.total_ops += th.DoubleTensor([matmul_ops])
276
+
277
+
278
+ class UNetModel(nn.Module):
279
+ """
280
+ The full UNet model with attention and timestep embedding.
281
+
282
+ :param in_channels: channels in the input Tensor.
283
+ :param model_channels: base channel count for the model.
284
+ :param out_channels: channels in the output Tensor.
285
+ :param num_res_blocks: number of residual blocks per downsample.
286
+ :param attention_resolutions: a collection of downsample rates at which
287
+ attention will take place. May be a set, list, or tuple.
288
+ For example, if this contains 4, then at 4x downsampling, attention
289
+ will be used.
290
+ :param dropout: the dropout probability.
291
+ :param channel_mult: channel multiplier for each level of the UNet.
292
+ :param conv_resample: if True, use learned convolutions for upsampling and
293
+ downsampling.
294
+ :param dims: determines if the signal is 1D, 2D, or 3D.
295
+ :param num_classes: if specified (as an int), then this model will be
296
+ class-conditional with `num_classes` classes.
297
+ :param use_checkpoint: use gradient checkpointing to reduce memory usage.
298
+ :param num_heads: the number of attention heads in each attention layer.
299
+ """
300
+
301
+ def __init__(
302
+ self,
303
+ in_channels,
304
+ model_channels,
305
+ out_channels,
306
+ num_res_blocks,
307
+ attention_resolutions,
308
+ dropout=0,
309
+ channel_mult=(1, 2, 4, 8),
310
+ conv_resample=True,
311
+ dims=2,
312
+ num_classes=None,
313
+ use_checkpoint=False,
314
+ num_heads=1,
315
+ num_heads_upsample=-1,
316
+ use_scale_shift_norm=False,
317
+ ):
318
+ super().__init__()
319
+
320
+ if num_heads_upsample == -1:
321
+ num_heads_upsample = num_heads
322
+
323
+ self.in_channels = in_channels
324
+ self.model_channels = model_channels
325
+ self.out_channels = out_channels
326
+ self.num_res_blocks = num_res_blocks
327
+ self.attention_resolutions = attention_resolutions
328
+ self.dropout = dropout
329
+ self.channel_mult = channel_mult
330
+ self.conv_resample = conv_resample
331
+ self.num_classes = num_classes
332
+ self.use_checkpoint = use_checkpoint
333
+ self.num_heads = num_heads
334
+ self.num_heads_upsample = num_heads_upsample
335
+
336
+ time_embed_dim = model_channels * 4
337
+ self.time_embed = nn.Sequential(
338
+ linear(model_channels, time_embed_dim),
339
+ SiLU(),
340
+ linear(time_embed_dim, time_embed_dim),
341
+ )
342
+
343
+ if self.num_classes is not None:
344
+ self.label_emb = nn.Embedding(num_classes, time_embed_dim)
345
+
346
+ self.input_blocks = nn.ModuleList(
347
+ [
348
+ TimestepEmbedSequential(
349
+ conv_nd(dims, in_channels, model_channels, 3, padding=1)
350
+ )
351
+ ]
352
+ )
353
+ input_block_chans = [model_channels]
354
+ ch = model_channels
355
+ ds = 1
356
+ for level, mult in enumerate(channel_mult):
357
+ for _ in range(num_res_blocks):
358
+ layers = [
359
+ ResBlock(
360
+ ch,
361
+ time_embed_dim,
362
+ dropout,
363
+ out_channels=mult * model_channels,
364
+ dims=dims,
365
+ use_checkpoint=use_checkpoint,
366
+ use_scale_shift_norm=use_scale_shift_norm,
367
+ )
368
+ ]
369
+ ch = mult * model_channels
370
+ if ds in attention_resolutions:
371
+ layers.append(
372
+ AttentionBlock(
373
+ ch, use_checkpoint=use_checkpoint, num_heads=num_heads
374
+ )
375
+ )
376
+ self.input_blocks.append(TimestepEmbedSequential(*layers))
377
+ input_block_chans.append(ch)
378
+ if level != len(channel_mult) - 1:
379
+ self.input_blocks.append(
380
+ TimestepEmbedSequential(Downsample(ch, conv_resample, dims=dims))
381
+ )
382
+ input_block_chans.append(ch)
383
+ ds *= 2
384
+
385
+ self.middle_block = TimestepEmbedSequential(
386
+ ResBlock(
387
+ ch,
388
+ time_embed_dim,
389
+ dropout,
390
+ dims=dims,
391
+ use_checkpoint=use_checkpoint,
392
+ use_scale_shift_norm=use_scale_shift_norm,
393
+ ),
394
+ AttentionBlock(ch, use_checkpoint=use_checkpoint, num_heads=num_heads),
395
+ ResBlock(
396
+ ch,
397
+ time_embed_dim,
398
+ dropout,
399
+ dims=dims,
400
+ use_checkpoint=use_checkpoint,
401
+ use_scale_shift_norm=use_scale_shift_norm,
402
+ ),
403
+ )
404
+
405
+ self.output_blocks = nn.ModuleList([])
406
+ for level, mult in list(enumerate(channel_mult))[::-1]:
407
+ for i in range(num_res_blocks + 1):
408
+ layers = [
409
+ ResBlock(
410
+ ch + input_block_chans.pop(),
411
+ time_embed_dim,
412
+ dropout,
413
+ out_channels=model_channels * mult,
414
+ dims=dims,
415
+ use_checkpoint=use_checkpoint,
416
+ use_scale_shift_norm=use_scale_shift_norm,
417
+ )
418
+ ]
419
+ ch = model_channels * mult
420
+ if ds in attention_resolutions:
421
+ layers.append(
422
+ AttentionBlock(
423
+ ch,
424
+ use_checkpoint=use_checkpoint,
425
+ num_heads=num_heads_upsample,
426
+ )
427
+ )
428
+ if level and i == num_res_blocks:
429
+ layers.append(Upsample(ch, conv_resample, dims=dims))
430
+ ds //= 2
431
+ self.output_blocks.append(TimestepEmbedSequential(*layers))
432
+
433
+ self.out = nn.Sequential(
434
+ normalization(ch),
435
+ SiLU(),
436
+ zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)),
437
+ )
438
+
439
+ def convert_to_fp16(self):
440
+ """
441
+ Convert the torso of the model to float16.
442
+ """
443
+ self.input_blocks.apply(convert_module_to_f16)
444
+ self.middle_block.apply(convert_module_to_f16)
445
+ self.output_blocks.apply(convert_module_to_f16)
446
+
447
+ def convert_to_fp32(self):
448
+ """
449
+ Convert the torso of the model to float32.
450
+ """
451
+ self.input_blocks.apply(convert_module_to_f32)
452
+ self.middle_block.apply(convert_module_to_f32)
453
+ self.output_blocks.apply(convert_module_to_f32)
454
+
455
+ @property
456
+ def inner_dtype(self):
457
+ """
458
+ Get the dtype used by the torso of the model.
459
+ """
460
+ return next(self.input_blocks.parameters()).dtype
461
+
462
+ def forward(self, x, timesteps, y=None):
463
+ """
464
+ Apply the model to an input batch.
465
+
466
+ :param x: an [N x C x ...] Tensor of inputs.
467
+ :param timesteps: a 1-D batch of timesteps.
468
+ :param y: an [N] Tensor of labels, if class-conditional.
469
+ :return: an [N x C x ...] Tensor of outputs.
470
+ """
471
+ assert (y is not None) == (
472
+ self.num_classes is not None
473
+ ), "must specify y if and only if the model is class-conditional"
474
+
475
+ hs = []
476
+ emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
477
+
478
+ if self.num_classes is not None:
479
+ assert y.shape == (x.shape[0],)
480
+ emb = emb + self.label_emb(y)
481
+
482
+ h = x.type(self.inner_dtype)
483
+ for module in self.input_blocks:
484
+ h = module(h, emb)
485
+ hs.append(h)
486
+ h = self.middle_block(h, emb)
487
+ for module in self.output_blocks:
488
+ cat_in = th.cat([h, hs.pop()], dim=1)
489
+ h = module(cat_in, emb)
490
+ h = h.type(x.dtype)
491
+ return self.out(h)
492
+
493
+ def get_feature_vectors(self, x, timesteps, y=None):
494
+ """
495
+ Apply the model and return all of the intermediate tensors.
496
+
497
+ :param x: an [N x C x ...] Tensor of inputs.
498
+ :param timesteps: a 1-D batch of timesteps.
499
+ :param y: an [N] Tensor of labels, if class-conditional.
500
+ :return: a dict with the following keys:
501
+ - 'down': a list of hidden state tensors from downsampling.
502
+ - 'middle': the tensor of the output of the lowest-resolution
503
+ block in the model.
504
+ - 'up': a list of hidden state tensors from upsampling.
505
+ """
506
+ hs = []
507
+ emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
508
+ if self.num_classes is not None:
509
+ assert y.shape == (x.shape[0],)
510
+ emb = emb + self.label_emb(y)
511
+ result = dict(down=[], up=[])
512
+ h = x.type(self.inner_dtype)
513
+ for module in self.input_blocks:
514
+ h = module(h, emb)
515
+ hs.append(h)
516
+ result["down"].append(h.type(x.dtype))
517
+ h = self.middle_block(h, emb)
518
+ result["middle"] = h.type(x.dtype)
519
+ for module in self.output_blocks:
520
+ cat_in = th.cat([h, hs.pop()], dim=1)
521
+ h = module(cat_in, emb)
522
+ result["up"].append(h.type(x.dtype))
523
+ return result
524
+
525
+
526
+ class SuperResModel(UNetModel):
527
+ """
528
+ A UNetModel that performs super-resolution.
529
+
530
+ Expects an extra kwarg `low_res` to condition on a low-resolution image.
531
+ """
532
+
533
+ def __init__(self, in_channels, *args, **kwargs):
534
+ super().__init__(in_channels * 2, *args, **kwargs)
535
+
536
+ def forward(self, x, timesteps, low_res=None, **kwargs):
537
+ _, _, new_height, new_width = x.shape
538
+ upsampled = F.interpolate(low_res, (new_height, new_width), mode="bilinear")
539
+ x = th.cat([x, upsampled], dim=1)
540
+ return super().forward(x, timesteps, **kwargs)
541
+
542
+ def get_feature_vectors(self, x, timesteps, low_res=None, **kwargs):
543
+ _, new_height, new_width, _ = x.shape
544
+ upsampled = F.interpolate(low_res, (new_height, new_width), mode="bilinear")
545
+ x = th.cat([x, upsampled], dim=1)
546
+ return super().get_feature_vectors(x, timesteps, **kwargs)
547
+
scripts/image_nll.py ADDED
@@ -0,0 +1,96 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 improved_diffusion import dist_util, logger
12
+ from improved_diffusion.image_datasets import load_data
13
+ from improved_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()
scripts/image_sample.py ADDED
@@ -0,0 +1,106 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 improved_diffusion import dist_util, logger
14
+ from improved_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
+ model.eval()
38
+
39
+ logger.log("sampling...")
40
+ all_images = []
41
+ all_labels = []
42
+ while len(all_images) * args.batch_size < args.num_samples:
43
+ model_kwargs = {}
44
+ if args.class_cond:
45
+ classes = th.randint(
46
+ low=0, high=NUM_CLASSES, size=(args.batch_size,), device=dist_util.dev()
47
+ )
48
+ model_kwargs["y"] = classes
49
+ sample_fn = (
50
+ diffusion.p_sample_loop if not args.use_ddim else diffusion.ddim_sample_loop
51
+ )
52
+ sample = sample_fn(
53
+ model,
54
+ (args.batch_size, 3, args.image_size, args.image_size),
55
+ clip_denoised=args.clip_denoised,
56
+ model_kwargs=model_kwargs,
57
+ )
58
+ sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8)
59
+ sample = sample.permute(0, 2, 3, 1)
60
+ sample = sample.contiguous()
61
+
62
+ gathered_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())]
63
+ dist.all_gather(gathered_samples, sample) # gather not supported with NCCL
64
+ all_images.extend([sample.cpu().numpy() for sample in gathered_samples])
65
+ if args.class_cond:
66
+ gathered_labels = [
67
+ th.zeros_like(classes) for _ in range(dist.get_world_size())
68
+ ]
69
+ dist.all_gather(gathered_labels, classes)
70
+ all_labels.extend([labels.cpu().numpy() for labels in gathered_labels])
71
+ logger.log(f"created {len(all_images) * args.batch_size} samples")
72
+
73
+ arr = np.concatenate(all_images, axis=0)
74
+ arr = arr[: args.num_samples]
75
+ if args.class_cond:
76
+ label_arr = np.concatenate(all_labels, axis=0)
77
+ label_arr = label_arr[: args.num_samples]
78
+ if dist.get_rank() == 0:
79
+ shape_str = "x".join([str(x) for x in arr.shape])
80
+ out_path = os.path.join(logger.get_dir(), f"samples_{shape_str}.npz")
81
+ logger.log(f"saving to {out_path}")
82
+ if args.class_cond:
83
+ np.savez(out_path, arr, label_arr)
84
+ else:
85
+ np.savez(out_path, arr)
86
+
87
+ dist.barrier()
88
+ logger.log("sampling complete")
89
+
90
+
91
+ def create_argparser():
92
+ defaults = dict(
93
+ clip_denoised=True,
94
+ num_samples=10000,
95
+ batch_size=16,
96
+ use_ddim=False,
97
+ model_path="",
98
+ )
99
+ defaults.update(model_and_diffusion_defaults())
100
+ parser = argparse.ArgumentParser()
101
+ add_dict_to_argparser(parser, defaults)
102
+ return parser
103
+
104
+
105
+ if __name__ == "__main__":
106
+ main()
scripts/image_train.py ADDED
@@ -0,0 +1,83 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Train a diffusion model on images.
3
+ """
4
+
5
+ import argparse
6
+
7
+ from improved_diffusion import dist_util, logger
8
+ from improved_diffusion.image_datasets import load_data
9
+ from improved_diffusion.resample import create_named_schedule_sampler
10
+ from improved_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 improved_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
+
40
+ logger.log("training...")
41
+ TrainLoop(
42
+ model=model,
43
+ diffusion=diffusion,
44
+ data=data,
45
+ batch_size=args.batch_size,
46
+ microbatch=args.microbatch,
47
+ lr=args.lr,
48
+ ema_rate=args.ema_rate,
49
+ log_interval=args.log_interval,
50
+ save_interval=args.save_interval,
51
+ resume_checkpoint=args.resume_checkpoint,
52
+ use_fp16=args.use_fp16,
53
+ fp16_scale_growth=args.fp16_scale_growth,
54
+ schedule_sampler=schedule_sampler,
55
+ weight_decay=args.weight_decay,
56
+ lr_anneal_steps=args.lr_anneal_steps,
57
+ ).run_loop()
58
+
59
+
60
+ def create_argparser():
61
+ defaults = dict(
62
+ data_dir="",
63
+ schedule_sampler="uniform",
64
+ lr=1e-4,
65
+ weight_decay=0.0,
66
+ lr_anneal_steps=0,
67
+ batch_size=1,
68
+ microbatch=-1, # -1 disables microbatches
69
+ ema_rate="0.9999", # comma-separated list of EMA values
70
+ log_interval=10,
71
+ save_interval=10000,
72
+ resume_checkpoint="",
73
+ use_fp16=False,
74
+ fp16_scale_growth=1e-3,
75
+ )
76
+ defaults.update(model_and_diffusion_defaults())
77
+ parser = argparse.ArgumentParser()
78
+ add_dict_to_argparser(parser, defaults)
79
+ return parser
80
+
81
+
82
+ if __name__ == "__main__":
83
+ main()
scripts/super_res_sample.py ADDED
@@ -0,0 +1,117 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 improved_diffusion import dist_util, logger
15
+ from improved_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
+ model.eval()
38
+
39
+ logger.log("loading data...")
40
+ data = load_data_for_worker(args.base_samples, args.batch_size, args.class_cond)
41
+
42
+ logger.log("creating samples...")
43
+ all_images = []
44
+ while len(all_images) * args.batch_size < args.num_samples:
45
+ model_kwargs = next(data)
46
+ model_kwargs = {k: v.to(dist_util.dev()) for k, v in model_kwargs.items()}
47
+ sample = diffusion.p_sample_loop(
48
+ model,
49
+ (args.batch_size, 3, args.large_size, args.large_size),
50
+ clip_denoised=args.clip_denoised,
51
+ model_kwargs=model_kwargs,
52
+ )
53
+ sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8)
54
+ sample = sample.permute(0, 2, 3, 1)
55
+ sample = sample.contiguous()
56
+
57
+ all_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())]
58
+ dist.all_gather(all_samples, sample) # gather not supported with NCCL
59
+ for sample in all_samples:
60
+ all_images.append(sample.cpu().numpy())
61
+ logger.log(f"created {len(all_images) * args.batch_size} samples")
62
+
63
+ arr = np.concatenate(all_images, axis=0)
64
+ arr = arr[: args.num_samples]
65
+ if dist.get_rank() == 0:
66
+ shape_str = "x".join([str(x) for x in arr.shape])
67
+ out_path = os.path.join(logger.get_dir(), f"samples_{shape_str}.npz")
68
+ logger.log(f"saving to {out_path}")
69
+ np.savez(out_path, arr)
70
+
71
+ dist.barrier()
72
+ logger.log("sampling complete")
73
+
74
+
75
+ def load_data_for_worker(base_samples, batch_size, class_cond):
76
+ with bf.BlobFile(base_samples, "rb") as f:
77
+ obj = np.load(f)
78
+ image_arr = obj["arr_0"]
79
+ if class_cond:
80
+ label_arr = obj["arr_1"]
81
+ rank = dist.get_rank()
82
+ num_ranks = dist.get_world_size()
83
+ buffer = []
84
+ label_buffer = []
85
+ while True:
86
+ for i in range(rank, len(image_arr), num_ranks):
87
+ buffer.append(image_arr[i])
88
+ if class_cond:
89
+ label_buffer.append(label_arr[i])
90
+ if len(buffer) == batch_size:
91
+ batch = th.from_numpy(np.stack(buffer)).float()
92
+ batch = batch / 127.5 - 1.0
93
+ batch = batch.permute(0, 3, 1, 2)
94
+ res = dict(low_res=batch)
95
+ if class_cond:
96
+ res["y"] = th.from_numpy(np.stack(label_buffer))
97
+ yield res
98
+ buffer, label_buffer = [], []
99
+
100
+
101
+ def create_argparser():
102
+ defaults = dict(
103
+ clip_denoised=True,
104
+ num_samples=10000,
105
+ batch_size=16,
106
+ use_ddim=False,
107
+ base_samples="",
108
+ model_path="",
109
+ )
110
+ defaults.update(sr_model_and_diffusion_defaults())
111
+ parser = argparse.ArgumentParser()
112
+ add_dict_to_argparser(parser, defaults)
113
+ return parser
114
+
115
+
116
+ if __name__ == "__main__":
117
+ main()
scripts/super_res_train.py ADDED
@@ -0,0 +1,98 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Train a super-resolution model.
3
+ """
4
+
5
+ import argparse
6
+
7
+ import torch.nn.functional as F
8
+
9
+ from improved_diffusion import dist_util, logger
10
+ from improved_diffusion.image_datasets import load_data
11
+ from improved_diffusion.resample import create_named_schedule_sampler
12
+ from improved_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 improved_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()
setup.py ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ from setuptools import setup
2
+
3
+ setup(
4
+ name="improved-diffusion",
5
+ py_modules=["improved_diffusion"],
6
+ install_requires=["blobfile>=1.0.5", "torch", "tqdm"],
7
+ )