├── cm ├── __init__.py ├── losses.py ├── dist_util.py ├── nn.py ├── image_datasets.py ├── random_util.py ├── resample.py ├── fp16_util.py ├── script_util.py ├── logger.py ├── train_util.py ├── unet.py └── network.py ├── assets └── imagenet64.png ├── .gitignore ├── setup.py ├── LICENSE ├── scripts ├── fid_eval.sh ├── distill_pid_diffusion.sh ├── image_sampling.sh ├── image_sample.py ├── cm_train.py └── fid_evaluation.py └── README.md /cm/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Codebase for "Improved Denoising Diffusion Probabilistic Models". 3 | """ 4 | -------------------------------------------------------------------------------- /assets/imagenet64.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pantheon5100/pid_diffusion/HEAD/assets/imagenet64.png -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Compiled python modules. 2 | *.pyc 3 | /cifar10 4 | /.hypothesis 5 | 6 | # Byte-compiled 7 | _pycache__/ 8 | .cache/ 9 | .idea/ 10 | 11 | # Python egg metadata, regenerated from source files by setuptools. 12 | /*.egg-info 13 | .eggs/ 14 | 15 | # PyPI distribution artifacts. 16 | build/ 17 | dist/ 18 | 19 | # Tests 20 | .pytest_cache/ 21 | 22 | # Other 23 | *.DS_Store 24 | 25 | /experiment 26 | /model_zoo/ 27 | /evaluations/*.npz 28 | 29 | /evaluations/cifar_img_fid 30 | /evaluations/classify_image_graph_def.pb 31 | 32 | !/model_zoo/README.md 33 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | setup( 4 | name="consistency-models", 5 | py_modules=["cm", "evaluations"], 6 | install_requires=[ 7 | "blobfile>=1.0.5", 8 | "torch", 9 | "clean-fid", 10 | "tqdm", 11 | "numpy", 12 | "scipy", 13 | "pandas", 14 | "Cython", 15 | "piq==0.7.0", 16 | "joblib==0.14.0", 17 | "albumentations==0.4.3", 18 | "lmdb", 19 | "clip @ git+https://github.com/openai/CLIP.git", 20 | "flash-attn==0.2.8", 21 | "pillow", 22 | ], 23 | ) 24 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 Zak 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. 22 | -------------------------------------------------------------------------------- /scripts/fid_eval.sh: -------------------------------------------------------------------------------- 1 | 2 | ################################################################## 3 | # For CIFAR model evaluation 4 | ################################################################## 5 | 6 | EXP_PATH="./model_zoo/pid_cifar" 7 | 8 | mpirun -np 1 python ./scripts/fid_evaluation.py \ 9 | --training_mode one_shot_pinn_edm_edm_one_shot \ 10 | --fid_dataset cifar10 \ 11 | --exp_dir $EXP_PATH\ 12 | --batch_size 125 \ 13 | --sigma_max 80 \ 14 | --sigma_min 0.002 \ 15 | --s_churn 0 \ 16 | --steps 35 \ 17 | --sampler oneshot \ 18 | --attention_resolutions "2" \ 19 | --class_cond False \ 20 | --dropout 0.0 \ 21 | --image_size 32 \ 22 | --num_channels 128 \ 23 | --num_res_blocks 4 \ 24 | --num_samples 50000 \ 25 | --resblock_updown True \ 26 | --use_fp16 False \ 27 | --use_scale_shift_norm True \ 28 | --weight_schedule uniform \ 29 | --seed 0 30 | 31 | ################################################################## 32 | 33 | 34 | ################################################################## 35 | # For ImageNet model evaluation 36 | ################################################################## 37 | 38 | # EXP_PATH="./experiment/pid_imagenet" 39 | 40 | # mpirun -np 1 python ./scripts/fid_evaluation.py \ 41 | # --training_mode one_shot_pinn_edm_edm_one_shot \ 42 | # --fid_dataset imagenet \ 43 | # --exp_dir $EXP_PATH\ 44 | # --batch_size 250 \ 45 | # --sigma_max 80 \ 46 | # --sigma_min 0.002 \ 47 | # --s_churn 0 \ 48 | # --sampler oneshot \ 49 | # --attention_resolutions 32,16,8 \ 50 | # --class_cond True \ 51 | # --dropout 0.0 \ 52 | # --image_size 64 \ 53 | # --num_channels 192 \ 54 | # --num_head_channels 64 \ 55 | # --num_res_blocks 3 \ 56 | # --num_samples 50000 \ 57 | # --resblock_updown True \ 58 | # --use_fp16 True \ 59 | # --use_scale_shift_norm True \ 60 | # --weight_schedule uniform 61 | 62 | ################################################################## 63 | -------------------------------------------------------------------------------- /scripts/distill_pid_diffusion.sh: -------------------------------------------------------------------------------- 1 | ######################################################################### 2 | # PID on CIFAR10 3 | ######################################################################### 4 | 5 | 6 | OPENAI_LOGDIR=./experiment/pid_cifar10 mpirun -np 8 python ./scripts/cm_train.py \ 7 | --training_mode one_shot_pinn_edm_edm \ 8 | --target_ema_mode fixed \ 9 | --start_ema 0.5 \ 10 | --scale_mode fixed \ 11 | --start_scales 250 \ 12 | --total_training_steps 800000 \ 13 | --loss_norm lpips \ 14 | --lr_anneal_steps 0 \ 15 | --teacher_model_path ./model_zoo/edm-cifar10-32x32-uncond-vp.ckpt \ 16 | --attention_resolutions "2" \ 17 | --class_cond False \ 18 | --use_scale_shift_norm True \ 19 | --dropout 0.0 \ 20 | --teacher_dropout 0.0 \ 21 | --ema_rate 0.999,0.9999,0.99995 \ 22 | --global_batch_size 512 \ 23 | --microbatch -1 \ 24 | --image_size 32 \ 25 | --lr 0.0002 \ 26 | --num_channels 128 \ 27 | --num_res_blocks 4 \ 28 | --resblock_updown True \ 29 | --schedule_sampler uniform \ 30 | --use_fp16 False \ 31 | --weight_decay 0.0 \ 32 | --weight_schedule uniform \ 33 | --optimizer radam 34 | 35 | ######################################################################### 36 | 37 | 38 | ######################################################################### 39 | # PID on ImageNet 64x64 40 | ######################################################################### 41 | 42 | # OPENAI_LOGDIR=./experiment/pid_imagenet mpirun -np 8 python ./scripts/cm_train.py \ 43 | # --training_mode one_shot_pinn_edm_edm \ 44 | # --target_ema_mode fixed \ 45 | # --start_ema 0.5 \ 46 | # --scale_mode fixed \ 47 | # --start_scales 250 \ 48 | # --total_training_steps 400000 \ 49 | # --loss_norm lpips \ 50 | # --lr_anneal_steps 0 \ 51 | # --teacher_model_path ./model_zoo/edm-imagenet-64x64-cond-adm.ckpt \ 52 | # --attention_resolutions "2" \ 53 | # --class_cond True \ 54 | # --use_scale_shift_norm True \ 55 | # --dropout 0.0 \ 56 | # --teacher_dropout 0.0 \ 57 | # --ema_rate 0.999,0.9999,0.99995 \ 58 | # --global_batch_size 2048 \ 59 | # --microbatch 32 \ 60 | # --image_size 64 \ 61 | # --lr 0.0001 \ 62 | # --num_channels 192 \ 63 | # --num_res_blocks 3 \ 64 | # --resblock_updown True \ 65 | # --schedule_sampler uniform \ 66 | # --use_fp16 True \ 67 | # --weight_decay 0.0 \ 68 | # --weight_schedule uniform \ 69 | # --optimizer radam 70 | 71 | ######################################################################### 72 | 73 | -------------------------------------------------------------------------------- /cm/losses.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers for various likelihood-based losses. These are ported from the original 3 | Ho et al. diffusion models codebase: 4 | https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/utils.py 5 | """ 6 | 7 | import numpy as np 8 | 9 | import torch as th 10 | 11 | 12 | def normal_kl(mean1, logvar1, mean2, logvar2): 13 | """ 14 | Compute the KL divergence between two gaussians. 15 | 16 | Shapes are automatically broadcasted, so batches can be compared to 17 | scalars, among other use cases. 18 | """ 19 | tensor = None 20 | for obj in (mean1, logvar1, mean2, logvar2): 21 | if isinstance(obj, th.Tensor): 22 | tensor = obj 23 | break 24 | assert tensor is not None, "at least one argument must be a Tensor" 25 | 26 | # Force variances to be Tensors. Broadcasting helps convert scalars to 27 | # Tensors, but it does not work for th.exp(). 28 | logvar1, logvar2 = [ 29 | x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor) 30 | for x in (logvar1, logvar2) 31 | ] 32 | 33 | return 0.5 * ( 34 | -1.0 35 | + logvar2 36 | - logvar1 37 | + th.exp(logvar1 - logvar2) 38 | + ((mean1 - mean2) ** 2) * th.exp(-logvar2) 39 | ) 40 | 41 | 42 | def approx_standard_normal_cdf(x): 43 | """ 44 | A fast approximation of the cumulative distribution function of the 45 | standard normal. 46 | """ 47 | return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3)))) 48 | 49 | 50 | def discretized_gaussian_log_likelihood(x, *, means, log_scales): 51 | """ 52 | Compute the log-likelihood of a Gaussian distribution discretizing to a 53 | given image. 54 | 55 | :param x: the target images. It is assumed that this was uint8 values, 56 | rescaled to the range [-1, 1]. 57 | :param means: the Gaussian mean Tensor. 58 | :param log_scales: the Gaussian log stddev Tensor. 59 | :return: a tensor like x of log probabilities (in nats). 60 | """ 61 | assert x.shape == means.shape == log_scales.shape 62 | centered_x = x - means 63 | inv_stdv = th.exp(-log_scales) 64 | plus_in = inv_stdv * (centered_x + 1.0 / 255.0) 65 | cdf_plus = approx_standard_normal_cdf(plus_in) 66 | min_in = inv_stdv * (centered_x - 1.0 / 255.0) 67 | cdf_min = approx_standard_normal_cdf(min_in) 68 | log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12)) 69 | log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12)) 70 | cdf_delta = cdf_plus - cdf_min 71 | log_probs = th.where( 72 | x < -0.999, 73 | log_cdf_plus, 74 | th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))), 75 | ) 76 | assert log_probs.shape == x.shape 77 | return log_probs 78 | -------------------------------------------------------------------------------- /cm/dist_util.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers for distributed training. 3 | """ 4 | 5 | import io 6 | import os 7 | import socket 8 | 9 | import blobfile as bf 10 | from mpi4py import MPI 11 | import torch as th 12 | import torch.distributed as dist 13 | 14 | # Change this to reflect your cluster layout. 15 | # The GPU for a given rank is (rank % GPUS_PER_NODE). 16 | GPUS_PER_NODE = 8 17 | 18 | SETUP_RETRY_COUNT = 3 19 | 20 | 21 | def setup_dist(): 22 | """ 23 | Setup a distributed process group. 24 | """ 25 | if dist.is_initialized(): 26 | return 27 | 28 | os.environ["CUDA_VISIBLE_DEVICES"] = f"{int(MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE)}," 29 | 30 | 31 | comm = MPI.COMM_WORLD 32 | backend = "gloo" if not th.cuda.is_available() else "nccl" 33 | 34 | if backend == "gloo": 35 | hostname = "localhost" 36 | else: 37 | hostname = socket.gethostbyname(socket.getfqdn()) 38 | os.environ["MASTER_ADDR"] = comm.bcast(hostname, root=0) 39 | os.environ["RANK"] = str(comm.rank) 40 | os.environ["WORLD_SIZE"] = str(comm.size) 41 | 42 | port = comm.bcast(_find_free_port(), root=0) 43 | os.environ["MASTER_PORT"] = str(port) 44 | dist.init_process_group(backend=backend, init_method="env://") 45 | 46 | 47 | 48 | def dev(): 49 | """ 50 | Get the device to use for torch.distributed. 51 | """ 52 | if th.cuda.is_available(): 53 | return th.device("cuda") 54 | return th.device("cpu") 55 | 56 | 57 | def load_state_dict(path, **kwargs): 58 | """ 59 | Load a PyTorch file without redundant fetches across MPI ranks. 60 | """ 61 | chunk_size = 2**30 # MPI has a relatively small size limit 62 | with bf.BlobFile(path, "rb") as f: 63 | data = f.read() 64 | # if MPI.COMM_WORLD.Get_rank() == 0: 65 | # with bf.BlobFile(path, "rb") as f: 66 | # data = f.read() 67 | # num_chunks = len(data) // chunk_size 68 | # if len(data) % chunk_size: 69 | # num_chunks += 1 70 | # MPI.COMM_WORLD.bcast(num_chunks) 71 | # for i in range(0, len(data), chunk_size): 72 | # MPI.COMM_WORLD.bcast(data[i : i + chunk_size]) 73 | # else: 74 | # num_chunks = MPI.COMM_WORLD.bcast(None) 75 | # data = bytes() 76 | # for _ in range(num_chunks): 77 | # data += MPI.COMM_WORLD.bcast(None) 78 | 79 | return th.load(io.BytesIO(data), **kwargs) 80 | 81 | 82 | def sync_params(params): 83 | """ 84 | Synchronize a sequence of Tensors across ranks from rank 0. 85 | """ 86 | for p in params: 87 | with th.no_grad(): 88 | dist.broadcast(p, 0) 89 | 90 | 91 | def _find_free_port(): 92 | try: 93 | s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 94 | s.bind(("", 0)) 95 | s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) 96 | return s.getsockname()[1] 97 | finally: 98 | s.close() 99 | -------------------------------------------------------------------------------- /scripts/image_sampling.sh: -------------------------------------------------------------------------------- 1 | ######################################################################### 2 | # Image generation for ImageNet edmedm model 3 | ######################################################################### 4 | 5 | # OPENAI_LOGDIR=./experiment/image_sampling/ImageNetEDM mpirun -np 1 python ./scripts/image_sample.py \ 6 | # --training_mode one_shot_pinn_edm_edm_teacher \ 7 | # --batch_size 128 \ 8 | # --sigma_max 80 \ 9 | # --sigma_min 0.002 \ 10 | # --s_churn 0 \ 11 | # --steps 79 \ 12 | # --sampler heun_deter \ 13 | # --model_path ./model_zoo/edm-imagenet-64x64-cond-adm.ckpt \ 14 | # --attention_resolutions 32,16,8 \ 15 | # --class_cond True \ 16 | # --dropout 0.1 \ 17 | # --image_size 64 \ 18 | # --num_channels 192 \ 19 | # --num_head_channels 64 \ 20 | # --num_res_blocks 3 \ 21 | # --num_samples 128 \ 22 | # --resblock_updown True \ 23 | # --use_fp16 False \ 24 | # --use_scale_shift_norm True \ 25 | # --weight_schedule uniform 26 | ######################################################################### 27 | 28 | ######################################################################### 29 | # Image generation for ImageNet oneshot model 30 | ######################################################################### 31 | 32 | # OPENAI_LOGDIR=./experiment/image_sampling/ImageNetPID mpirun -np 1 python ./scripts/image_sample.py \ 33 | # --training_mode one_shot_pinn_edm_edm_one_shot \ 34 | # --batch_size 128 \ 35 | # --sigma_max 80 \ 36 | # --sigma_min 0.002 \ 37 | # --s_churn 0 \ 38 | # --steps 79 \ 39 | # --sampler oneshot \ 40 | # --model_path ./model_zoo/pid_imagenet64.ckpt \ 41 | # --attention_resolutions 32,16,8 \ 42 | # --class_cond True \ 43 | # --dropout 0.1 \ 44 | # --image_size 64 \ 45 | # --num_channels 192 \ 46 | # --num_head_channels 64 \ 47 | # --num_res_blocks 3 \ 48 | # --num_samples 128 \ 49 | # --resblock_updown True \ 50 | # --use_fp16 False \ 51 | # --use_scale_shift_norm True \ 52 | # --weight_schedule uniform 53 | 54 | ######################################################################### 55 | 56 | 57 | ######################################################################### 58 | # Image generation for cifar edmedm model 59 | ######################################################################### 60 | 61 | OPENAI_LOGDIR=./experiment/image_sampling/CIFAREDM mpirun -np 1 python ./scripts/image_sample.py \ 62 | --training_mode one_shot_pinn_edm_edm_teacher \ 63 | --batch_size 128 \ 64 | --sigma_max 80 \ 65 | --sigma_min 0.002 \ 66 | --s_churn 0 \ 67 | --steps 35 \ 68 | --sampler heun_deter \ 69 | --model_path ./model_zoo/edm-cifar10-32x32-uncond-vp.ckpt \ 70 | --attention_resolutions "2" \ 71 | --class_cond False \ 72 | --dropout 0.0 \ 73 | --image_size 32 \ 74 | --num_channels 192 \ 75 | --num_channels 128 \ 76 | --num_res_blocks 4 \ 77 | --num_samples 128 \ 78 | --resblock_updown True \ 79 | --use_fp16 False \ 80 | --use_scale_shift_norm True \ 81 | --weight_schedule uniform 82 | ######################################################################### 83 | 84 | 85 | ######################################################################### 86 | # Image generation for cifar oneshot model 87 | ######################################################################### 88 | 89 | # OPENAI_LOGDIR=./experiment/image_sampling/CIFARPID mpirun -np 1 python ./scripts/image_sample.py \ 90 | # --training_mode one_shot_pinn_edm_edm_one_shot \ 91 | # --batch_size 128 \ 92 | # --sigma_max 80 \ 93 | # --sigma_min 0.002 \ 94 | # --s_churn 0 \ 95 | # --steps 35 \ 96 | # --sampler oneshot \ 97 | # --model_path ./model_zoo/pid_cifar.pt \ 98 | # --attention_resolutions "2" \ 99 | # --class_cond False \ 100 | # --dropout 0.0 \ 101 | # --image_size 32 \ 102 | # --num_channels 192 \ 103 | # --num_channels 128 \ 104 | # --num_res_blocks 4 \ 105 | # --num_samples 128 \ 106 | # --resblock_updown True \ 107 | # --use_fp16 False \ 108 | # --use_scale_shift_norm True \ 109 | # --weight_schedule uniform 110 | ######################################################################### 111 | 112 | -------------------------------------------------------------------------------- /scripts/image_sample.py: -------------------------------------------------------------------------------- 1 | """ 2 | Generate a large batch of image samples from a model and save them as a large 3 | numpy array. This can be used to produce samples for FID evaluation. 4 | """ 5 | 6 | import argparse 7 | import os 8 | 9 | import numpy as np 10 | import torch as th 11 | import torch.distributed as dist 12 | 13 | from cm import dist_util, logger 14 | from cm.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 | create_one_shot_edmedm_model_and_diffusion, 21 | ) 22 | from cm.random_util import get_generator 23 | from cm.karras_diffusion import karras_sample 24 | import copy 25 | import torchvision 26 | 27 | def main(): 28 | args = create_argparser().parse_args() 29 | 30 | dist_util.setup_dist() 31 | logger.configure() 32 | 33 | logger.log("creating model and diffusion...") 34 | model_and_diffusion_kwargs = args_to_dict(args, model_and_diffusion_defaults().keys()) 35 | 36 | if "one_shot_pinn_edm_edm_teacher" == args.training_mode: 37 | model, diffusion = create_one_shot_edmedm_model_and_diffusion(teacher_precond=True, 38 | **args_to_dict(args, model_and_diffusion_defaults().keys()), 39 | ) 40 | load_weights = dist_util.load_state_dict(args.model_path, map_location="cpu") 41 | new_state_dict = {} 42 | for k, v in load_weights.items(): 43 | if "map_augment" in k: 44 | continue 45 | new_key = k.replace("model.", "") 46 | new_state_dict[new_key] = v 47 | 48 | model.load_state_dict(new_state_dict) 49 | elif "one_shot_pinn_edm_edm_one_shot" == args.training_mode: 50 | model, diffusion = create_one_shot_edmedm_model_and_diffusion(teacher_precond=False, 51 | **args_to_dict(args, model_and_diffusion_defaults().keys()), 52 | ) 53 | load_weights = dist_util.load_state_dict(args.model_path, map_location="cpu") 54 | new_state_dict = {} 55 | for k, v in load_weights.items(): 56 | if "map_augment" in k: 57 | continue 58 | new_key = k.replace("model.", "") 59 | new_state_dict[new_key] = v 60 | 61 | model.load_state_dict(new_state_dict) 62 | 63 | 64 | model.to(dist_util.dev()) 65 | if args.use_fp16: 66 | model.convert_to_fp16() 67 | model.eval() 68 | 69 | logger.log("sampling...") 70 | if args.sampler == "multistep": 71 | assert len(args.ts) > 0 72 | ts = tuple(int(x) for x in args.ts.split(",")) 73 | else: 74 | ts = None 75 | 76 | all_images = [] 77 | all_labels = [] 78 | generator = get_generator(args.generator, args.num_samples, args.seed) 79 | 80 | while len(all_images) * args.batch_size < args.num_samples: 81 | model_kwargs = {} 82 | if args.class_cond: 83 | classes = th.randint( 84 | low=0, high=NUM_CLASSES, size=(args.batch_size,), device=dist_util.dev() 85 | ) 86 | model_kwargs["y"] = classes 87 | 88 | sample_ori = karras_sample( 89 | diffusion, 90 | model, 91 | (args.batch_size, 3, args.image_size, args.image_size), 92 | steps=args.steps, 93 | model_kwargs=model_kwargs, 94 | device=dist_util.dev(), 95 | clip_denoised=args.clip_denoised, 96 | sampler=args.sampler, 97 | sigma_min=args.sigma_min, 98 | sigma_max=args.sigma_max, 99 | s_churn=args.s_churn, 100 | s_tmin=args.s_tmin, 101 | s_tmax=args.s_tmax, 102 | s_noise=args.s_noise, 103 | generator=generator, 104 | ts=ts, 105 | ) 106 | sample = ((sample_ori + 1) * 127.5).clamp(0, 255).to(th.uint8) 107 | sample = sample.permute(0, 2, 3, 1) 108 | sample = sample.contiguous() 109 | 110 | gathered_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())] 111 | dist.all_gather(gathered_samples, sample) # gather not supported with NCCL 112 | all_images.extend([sample.cpu().numpy() for sample in gathered_samples]) 113 | if args.class_cond: 114 | gathered_labels = [ 115 | th.zeros_like(classes) for _ in range(dist.get_world_size()) 116 | ] 117 | dist.all_gather(gathered_labels, classes) 118 | all_labels.extend([labels.cpu().numpy() for labels in gathered_labels]) 119 | logger.log(f"created {len(all_images) * args.batch_size} samples") 120 | 121 | if len(all_images) <= args.batch_size and dist.get_rank() == 0: 122 | torchvision.utils.save_image((sample_ori + 1.)/2., os.path.join(logger.get_dir(), f"samples_first_batch.png")) 123 | 124 | arr = np.concatenate(all_images, axis=0) 125 | arr = arr[: args.num_samples] 126 | if args.class_cond: 127 | label_arr = np.concatenate(all_labels, axis=0) 128 | label_arr = label_arr[: args.num_samples] 129 | if dist.get_rank() == 0: 130 | shape_str = "x".join([str(x) for x in arr.shape]) 131 | out_path = os.path.join(logger.get_dir(), f"samples_{shape_str}.npz") 132 | logger.log(f"saving to {out_path}") 133 | 134 | if args.class_cond: 135 | np.savez(out_path, arr, label_arr) 136 | else: 137 | np.savez(out_path, arr) 138 | 139 | dist.barrier() 140 | logger.log("sampling complete") 141 | 142 | 143 | def create_argparser(): 144 | defaults = dict( 145 | training_mode="edm", 146 | generator="determ", 147 | clip_denoised=True, 148 | num_samples=10000, 149 | batch_size=16, 150 | sampler="heun", 151 | s_churn=0.0, 152 | s_tmin=0.0, 153 | s_tmax=float("inf"), 154 | s_noise=1.0, 155 | steps=40, 156 | model_path="", 157 | seed=42, 158 | ts="", 159 | ) 160 | defaults.update(model_and_diffusion_defaults()) 161 | parser = argparse.ArgumentParser() 162 | add_dict_to_argparser(parser, defaults) 163 | return parser 164 | 165 | 166 | if __name__ == "__main__": 167 | main() 168 | -------------------------------------------------------------------------------- /cm/nn.py: -------------------------------------------------------------------------------- 1 | """ 2 | Various utilities for neural networks. 3 | """ 4 | 5 | import math 6 | 7 | import torch as th 8 | import torch.nn as nn 9 | import numpy as np 10 | import torch.nn.functional as F 11 | 12 | 13 | # PyTorch 1.7 has SiLU, but we support PyTorch 1.5. 14 | class SiLU(nn.Module): 15 | def forward(self, x): 16 | return x * th.sigmoid(x) 17 | 18 | 19 | class GroupNorm32(nn.GroupNorm): 20 | def forward(self, x): 21 | return super().forward(x.float()).type(x.dtype) 22 | 23 | 24 | def conv_nd(dims, *args, **kwargs): 25 | """ 26 | Create a 1D, 2D, or 3D convolution module. 27 | """ 28 | if dims == 1: 29 | return nn.Conv1d(*args, **kwargs) 30 | elif dims == 2: 31 | return nn.Conv2d(*args, **kwargs) 32 | elif dims == 3: 33 | return nn.Conv3d(*args, **kwargs) 34 | raise ValueError(f"unsupported dimensions: {dims}") 35 | 36 | 37 | def linear(*args, **kwargs): 38 | """ 39 | Create a linear module. 40 | """ 41 | return nn.Linear(*args, **kwargs) 42 | 43 | 44 | def avg_pool_nd(dims, *args, **kwargs): 45 | """ 46 | Create a 1D, 2D, or 3D average pooling module. 47 | """ 48 | if dims == 1: 49 | return nn.AvgPool1d(*args, **kwargs) 50 | elif dims == 2: 51 | return nn.AvgPool2d(*args, **kwargs) 52 | elif dims == 3: 53 | return nn.AvgPool3d(*args, **kwargs) 54 | raise ValueError(f"unsupported dimensions: {dims}") 55 | 56 | 57 | def update_ema(target_params, source_params, rate=0.99): 58 | """ 59 | Update target parameters to be closer to those of source parameters using 60 | an exponential moving average. 61 | 62 | :param target_params: the target parameter sequence. 63 | :param source_params: the source parameter sequence. 64 | :param rate: the EMA rate (closer to 1 means slower). 65 | """ 66 | for targ, src in zip(target_params, source_params): 67 | targ.detach().mul_(rate).add_(src, alpha=1 - rate) 68 | 69 | 70 | def zero_module(module): 71 | """ 72 | Zero out the parameters of a module and return it. 73 | """ 74 | for p in module.parameters(): 75 | p.detach().zero_() 76 | return module 77 | 78 | 79 | def scale_module(module, scale): 80 | """ 81 | Scale the parameters of a module and return it. 82 | """ 83 | for p in module.parameters(): 84 | p.detach().mul_(scale) 85 | return module 86 | 87 | 88 | def mean_flat(tensor): 89 | """ 90 | Take the mean over all non-batch dimensions. 91 | """ 92 | return tensor.mean(dim=list(range(1, len(tensor.shape)))) 93 | 94 | 95 | def append_dims(x, target_dims): 96 | """Appends dimensions to the end of a tensor until it has target_dims dimensions.""" 97 | dims_to_append = target_dims - x.ndim 98 | if dims_to_append < 0: 99 | raise ValueError( 100 | f"input has {x.ndim} dims but target_dims is {target_dims}, which is less" 101 | ) 102 | return x[(...,) + (None,) * dims_to_append] 103 | 104 | 105 | def append_zero(x): 106 | return th.cat([x, x.new_zeros([1])]) 107 | 108 | 109 | def normalization(channels): 110 | """ 111 | Make a standard normalization layer. 112 | 113 | :param channels: number of input channels. 114 | :return: an nn.Module for normalization. 115 | """ 116 | return GroupNorm32(32, channels) 117 | 118 | 119 | def timestep_embedding(timesteps, dim, max_period=10000): 120 | """ 121 | Create sinusoidal timestep embeddings. 122 | 123 | :param timesteps: a 1-D Tensor of N indices, one per batch element. 124 | These may be fractional. 125 | :param dim: the dimension of the output. 126 | :param max_period: controls the minimum frequency of the embeddings. 127 | :return: an [N x dim] Tensor of positional embeddings. 128 | """ 129 | half = dim // 2 130 | freqs = th.exp( 131 | -math.log(max_period) * th.arange(start=0, end=half, dtype=th.float32) / half 132 | ).to(device=timesteps.device) 133 | args = timesteps[:, None].float() * freqs[None] 134 | embedding = th.cat([th.cos(args), th.sin(args)], dim=-1) 135 | if dim % 2: 136 | embedding = th.cat([embedding, th.zeros_like(embedding[:, :1])], dim=-1) 137 | return embedding 138 | 139 | 140 | def checkpoint(func, inputs, params, flag): 141 | """ 142 | Evaluate a function without caching intermediate activations, allowing for 143 | reduced memory at the expense of extra compute in the backward pass. 144 | 145 | :param func: the function to evaluate. 146 | :param inputs: the argument sequence to pass to `func`. 147 | :param params: a sequence of parameters `func` depends on but does not 148 | explicitly take as arguments. 149 | :param flag: if False, disable gradient checkpointing. 150 | """ 151 | if flag: 152 | args = tuple(inputs) + tuple(params) 153 | return CheckpointFunction.apply(func, len(inputs), *args) 154 | else: 155 | return func(*inputs) 156 | 157 | 158 | class CheckpointFunction(th.autograd.Function): 159 | @staticmethod 160 | def forward(ctx, run_function, length, *args): 161 | ctx.run_function = run_function 162 | ctx.input_tensors = list(args[:length]) 163 | ctx.input_params = list(args[length:]) 164 | with th.no_grad(): 165 | output_tensors = ctx.run_function(*ctx.input_tensors) 166 | return output_tensors 167 | 168 | @staticmethod 169 | def backward(ctx, *output_grads): 170 | ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors] 171 | with th.enable_grad(): 172 | # Fixes a bug where the first op in run_function modifies the 173 | # Tensor storage in place, which is not allowed for detach()'d 174 | # Tensors. 175 | shallow_copies = [x.view_as(x) for x in ctx.input_tensors] 176 | output_tensors = ctx.run_function(*shallow_copies) 177 | input_grads = th.autograd.grad( 178 | output_tensors, 179 | ctx.input_tensors + ctx.input_params, 180 | output_grads, 181 | allow_unused=True, 182 | ) 183 | del ctx.input_tensors 184 | del ctx.input_params 185 | del output_tensors 186 | return (None, None) + input_grads 187 | -------------------------------------------------------------------------------- /cm/image_datasets.py: -------------------------------------------------------------------------------- 1 | import math 2 | import random 3 | 4 | from PIL import Image 5 | import blobfile as bf 6 | from mpi4py import MPI 7 | import numpy as np 8 | from torch.utils.data import DataLoader, Dataset 9 | 10 | 11 | def load_data( 12 | *, 13 | data_dir, 14 | batch_size, 15 | image_size, 16 | class_cond=False, 17 | deterministic=False, 18 | random_crop=False, 19 | random_flip=True, 20 | ): 21 | """ 22 | For a dataset, create a generator over (images, kwargs) pairs. 23 | 24 | Each images is an NCHW float tensor, and the kwargs dict contains zero or 25 | more keys, each of which map to a batched Tensor of their own. 26 | The kwargs dict can be used for class labels, in which case the key is "y" 27 | and the values are integer tensors of class labels. 28 | 29 | :param data_dir: a dataset directory. 30 | :param batch_size: the batch size of each returned pair. 31 | :param image_size: the size to which images are resized. 32 | :param class_cond: if True, include a "y" key in returned dicts for class 33 | label. If classes are not available and this is true, an 34 | exception will be raised. 35 | :param deterministic: if True, yield results in a deterministic order. 36 | :param random_crop: if True, randomly crop the images for augmentation. 37 | :param random_flip: if True, randomly flip the images for augmentation. 38 | """ 39 | if not data_dir: 40 | raise ValueError("unspecified data directory") 41 | all_files = _list_image_files_recursively(data_dir) 42 | classes = None 43 | if class_cond: 44 | # Assume classes are the first part of the filename, 45 | # before an underscore. 46 | class_names = [bf.basename(path).split("_")[0] for path in all_files] 47 | sorted_classes = {x: i for i, x in enumerate(sorted(set(class_names)))} 48 | classes = [sorted_classes[x] for x in class_names] 49 | dataset = ImageDataset( 50 | image_size, 51 | all_files, 52 | classes=classes, 53 | shard=MPI.COMM_WORLD.Get_rank(), 54 | num_shards=MPI.COMM_WORLD.Get_size(), 55 | random_crop=random_crop, 56 | random_flip=random_flip, 57 | ) 58 | if deterministic: 59 | loader = DataLoader( 60 | dataset, batch_size=batch_size, shuffle=False, num_workers=4, drop_last=True 61 | ) 62 | else: 63 | loader = DataLoader( 64 | dataset, batch_size=batch_size, shuffle=True, num_workers=4, drop_last=True 65 | ) 66 | while True: 67 | yield from loader 68 | 69 | 70 | def _list_image_files_recursively(data_dir): 71 | results = [] 72 | for entry in sorted(bf.listdir(data_dir)): 73 | full_path = bf.join(data_dir, entry) 74 | ext = entry.split(".")[-1] 75 | if "." in entry and ext.lower() in ["jpg", "jpeg", "png", "gif"]: 76 | results.append(full_path) 77 | elif bf.isdir(full_path): 78 | results.extend(_list_image_files_recursively(full_path)) 79 | return results 80 | 81 | 82 | class ImageDataset(Dataset): 83 | def __init__( 84 | self, 85 | resolution, 86 | image_paths, 87 | classes=None, 88 | shard=0, 89 | num_shards=1, 90 | random_crop=False, 91 | random_flip=True, 92 | ): 93 | super().__init__() 94 | self.resolution = resolution 95 | self.local_images = image_paths[shard:][::num_shards] 96 | self.local_classes = None if classes is None else classes[shard:][::num_shards] 97 | self.random_crop = random_crop 98 | self.random_flip = random_flip 99 | 100 | def __len__(self): 101 | return len(self.local_images) 102 | 103 | def __getitem__(self, idx): 104 | path = self.local_images[idx] 105 | with bf.BlobFile(path, "rb") as f: 106 | pil_image = Image.open(f) 107 | pil_image.load() 108 | pil_image = pil_image.convert("RGB") 109 | 110 | if self.random_crop: 111 | arr = random_crop_arr(pil_image, self.resolution) 112 | else: 113 | arr = center_crop_arr(pil_image, self.resolution) 114 | 115 | if self.random_flip and random.random() < 0.5: 116 | arr = arr[:, ::-1] 117 | 118 | arr = arr.astype(np.float32) / 127.5 - 1 119 | 120 | out_dict = {} 121 | if self.local_classes is not None: 122 | out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64) 123 | return np.transpose(arr, [2, 0, 1]), out_dict 124 | 125 | 126 | def center_crop_arr(pil_image, image_size): 127 | # We are not on a new enough PIL to support the `reducing_gap` 128 | # argument, which uses BOX downsampling at powers of two first. 129 | # Thus, we do it by hand to improve downsample quality. 130 | while min(*pil_image.size) >= 2 * image_size: 131 | pil_image = pil_image.resize( 132 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX 133 | ) 134 | 135 | scale = image_size / min(*pil_image.size) 136 | pil_image = pil_image.resize( 137 | tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC 138 | ) 139 | 140 | arr = np.array(pil_image) 141 | crop_y = (arr.shape[0] - image_size) // 2 142 | crop_x = (arr.shape[1] - image_size) // 2 143 | return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] 144 | 145 | 146 | def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0): 147 | min_smaller_dim_size = math.ceil(image_size / max_crop_frac) 148 | max_smaller_dim_size = math.ceil(image_size / min_crop_frac) 149 | smaller_dim_size = random.randrange(min_smaller_dim_size, max_smaller_dim_size + 1) 150 | 151 | # We are not on a new enough PIL to support the `reducing_gap` 152 | # argument, which uses BOX downsampling at powers of two first. 153 | # Thus, we do it by hand to improve downsample quality. 154 | while min(*pil_image.size) >= 2 * smaller_dim_size: 155 | pil_image = pil_image.resize( 156 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX 157 | ) 158 | 159 | scale = smaller_dim_size / min(*pil_image.size) 160 | pil_image = pil_image.resize( 161 | tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC 162 | ) 163 | 164 | arr = np.array(pil_image) 165 | crop_y = random.randrange(arr.shape[0] - image_size + 1) 166 | crop_x = random.randrange(arr.shape[1] - image_size + 1) 167 | return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] 168 | -------------------------------------------------------------------------------- /cm/random_util.py: -------------------------------------------------------------------------------- 1 | import torch as th 2 | import torch.distributed as dist 3 | from . import dist_util 4 | 5 | 6 | def get_generator(generator, num_samples=0, seed=0): 7 | if generator == "dummy": 8 | return DummyGenerator() 9 | elif generator == "determ": 10 | return DeterministicGenerator(num_samples, seed) 11 | elif generator == "determ-indiv": 12 | return DeterministicIndividualGenerator(num_samples, seed) 13 | else: 14 | raise NotImplementedError 15 | 16 | 17 | class DummyGenerator: 18 | def randn(self, *args, **kwargs): 19 | return th.randn(*args, **kwargs) 20 | 21 | def randint(self, *args, **kwargs): 22 | return th.randint(*args, **kwargs) 23 | 24 | def randn_like(self, *args, **kwargs): 25 | return th.randn_like(*args, **kwargs) 26 | 27 | 28 | class DeterministicGenerator: 29 | """ 30 | RNG to deterministically sample num_samples samples that does not depend on batch_size or mpi_machines 31 | Uses a single rng and samples num_samples sized randomness and subsamples the current indices 32 | """ 33 | 34 | def __init__(self, num_samples, seed=0): 35 | if dist.is_initialized(): 36 | self.rank = dist.get_rank() 37 | self.world_size = dist.get_world_size() 38 | else: 39 | print("Warning: Distributed not initialised, using single rank") 40 | self.rank = 0 41 | self.world_size = 1 42 | self.num_samples = num_samples 43 | self.done_samples = 0 44 | self.seed = seed 45 | self.rng_cpu = th.Generator() 46 | if th.cuda.is_available(): 47 | self.rng_cuda = th.Generator(dist_util.dev()) 48 | self.set_seed(seed) 49 | 50 | def get_global_size_and_indices(self, size): 51 | global_size = (self.num_samples, *size[1:]) 52 | indices = th.arange( 53 | self.done_samples + self.rank, 54 | self.done_samples + self.world_size * int(size[0]), 55 | self.world_size, 56 | ) 57 | indices = th.clamp(indices, 0, self.num_samples - 1) 58 | assert ( 59 | len(indices) == size[0] 60 | ), f"rank={self.rank}, ws={self.world_size}, l={len(indices)}, bs={size[0]}" 61 | return global_size, indices 62 | 63 | def get_generator(self, device): 64 | return self.rng_cpu if th.device(device).type == "cpu" else self.rng_cuda 65 | 66 | def randn(self, *size, dtype=th.float, device="cpu"): 67 | global_size, indices = self.get_global_size_and_indices(size) 68 | generator = self.get_generator(device) 69 | return th.randn(*global_size, generator=generator, dtype=dtype, device=device)[ 70 | indices 71 | ] 72 | 73 | def randint(self, low, high, size, dtype=th.long, device="cpu"): 74 | global_size, indices = self.get_global_size_and_indices(size) 75 | generator = self.get_generator(device) 76 | return th.randint( 77 | low, high, generator=generator, size=global_size, dtype=dtype, device=device 78 | )[indices] 79 | 80 | def randn_like(self, tensor): 81 | size, dtype, device = tensor.size(), tensor.dtype, tensor.device 82 | return self.randn(*size, dtype=dtype, device=device) 83 | 84 | def set_done_samples(self, done_samples): 85 | self.done_samples = done_samples 86 | self.set_seed(self.seed) 87 | 88 | def get_seed(self): 89 | return self.seed 90 | 91 | def set_seed(self, seed): 92 | self.rng_cpu.manual_seed(seed) 93 | if th.cuda.is_available(): 94 | self.rng_cuda.manual_seed(seed) 95 | 96 | 97 | class DeterministicIndividualGenerator: 98 | """ 99 | RNG to deterministically sample num_samples samples that does not depend on batch_size or mpi_machines 100 | Uses a separate rng for each sample to reduce memoery usage 101 | """ 102 | 103 | def __init__(self, num_samples, seed=0): 104 | if dist.is_initialized(): 105 | self.rank = dist.get_rank() 106 | self.world_size = dist.get_world_size() 107 | else: 108 | print("Warning: Distributed not initialised, using single rank") 109 | self.rank = 0 110 | self.world_size = 1 111 | self.num_samples = num_samples 112 | self.done_samples = 0 113 | self.seed = seed 114 | self.rng_cpu = [th.Generator() for _ in range(num_samples)] 115 | if th.cuda.is_available(): 116 | self.rng_cuda = [th.Generator(dist_util.dev()) for _ in range(num_samples)] 117 | self.set_seed(seed) 118 | 119 | def get_size_and_indices(self, size): 120 | indices = th.arange( 121 | self.done_samples + self.rank, 122 | self.done_samples + self.world_size * int(size[0]), 123 | self.world_size, 124 | ) 125 | indices = th.clamp(indices, 0, self.num_samples - 1) 126 | assert ( 127 | len(indices) == size[0] 128 | ), f"rank={self.rank}, ws={self.world_size}, l={len(indices)}, bs={size[0]}" 129 | return (1, *size[1:]), indices 130 | 131 | def get_generator(self, device): 132 | return self.rng_cpu if th.device(device).type == "cpu" else self.rng_cuda 133 | 134 | def randn(self, *size, dtype=th.float, device="cpu"): 135 | size, indices = self.get_size_and_indices(size) 136 | generator = self.get_generator(device) 137 | return th.cat( 138 | [ 139 | th.randn(*size, generator=generator[i], dtype=dtype, device=device) 140 | for i in indices 141 | ], 142 | dim=0, 143 | ) 144 | 145 | def randint(self, low, high, size, dtype=th.long, device="cpu"): 146 | size, indices = self.get_size_and_indices(size) 147 | generator = self.get_generator(device) 148 | return th.cat( 149 | [ 150 | th.randint( 151 | low, 152 | high, 153 | generator=generator[i], 154 | size=size, 155 | dtype=dtype, 156 | device=device, 157 | ) 158 | for i in indices 159 | ], 160 | dim=0, 161 | ) 162 | 163 | def randn_like(self, tensor): 164 | size, dtype, device = tensor.size(), tensor.dtype, tensor.device 165 | return self.randn(*size, dtype=dtype, device=device) 166 | 167 | def set_done_samples(self, done_samples): 168 | self.done_samples = done_samples 169 | 170 | def get_seed(self): 171 | return self.seed 172 | 173 | def set_seed(self, seed): 174 | [ 175 | rng_cpu.manual_seed(i + self.num_samples * seed) 176 | for i, rng_cpu in enumerate(self.rng_cpu) 177 | ] 178 | if th.cuda.is_available(): 179 | [ 180 | rng_cuda.manual_seed(i + self.num_samples * seed) 181 | for i, rng_cuda in enumerate(self.rng_cuda) 182 | ] 183 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Physics Informed Distillation for Diffusion Models 2 | 3 | Joshua Tian Jin Tee*, Kang Zhang*, Hee Suk Yoon, Dhananjaya Nagaraja Gowda, Chanwoo Kim, Chang D. Yoo (*Equal contribution) 4 | 5 | This repository is the official implementation of the paper: [Physics Informed Distillation for Diffusion Models](https://openreview.net/forum?id=rOvaUsF996), accepted by Transactions on Machine Learning Research (TMLR). 6 | 7 | Diffusion models have recently emerged as a potent tool in generative modeling. However, their inherent iterative nature often results in sluggish image generation due to the requirement for multiple model evaluations. Recent progress has unveiled the intrinsic link between diffusion models and Probability Flow Ordinary Differential Equations (ODEs), thus enabling us to conceptualize diffusion models as ODE systems. Simultaneously, Physics Informed Neural Networks (PINNs) have substantiated their effectiveness in solving intricate differential equations through implicit modeling of their solutions. Building upon these foundational insights, we introduce Physics Informed Distillation (PID), which employs a student model to represent the solution of the ODE system corresponding to the teacher diffusion model, akin to the principles employed in PINNs. Through experiments on CIFAR 10 and ImageNet 64x64, we observe that PID achieves performance comparable to recent distillation methods. Notably, it demonstrates predictable trends concerning method-specific hyperparameters and eliminates the need for synthetic dataset generation during the distillation process. Both of which contribute to its easy-to-use nature as a distillation approach for Diffusion Models. 8 | 9 | 10 | 11 | Overview 12 | 13 | 14 | An overview of the proposed method, which involves training a model $\mathbf{x}_{\theta}(\mathbf{z}, \cdot )$ to approximate the true trajectory $\mathbf{x}(\mathbf{z}, \cdot )$. 15 | 16 | 17 | 18 | # 🔧 Environment Setup 19 | 20 | To install all packages in this codebase along with their dependencies, run 21 | ```sh 22 | conda create -n pid-diffusion python=3.9 23 | conda activate pid-diffusion 24 | conda install pytorch=1.13.1 torchvision=0.14.1 pytorch-cuda=11.6 -c pytorch -c nvidia 25 | conda install -c "nvidia/label/cuda-11.6.1" libcusolver-dev 26 | conda install mpi4py 27 | git clone https://github.com/pantheon5100/pid_diffusion.git 28 | cd pid_diffusion 29 | pip install -e . 30 | ``` 31 | 32 | # ⚡ Get Started 33 | 34 | ## 1. Preparing Pretrained Checkpoints 35 | ### Teacher Models for Distillation 36 | For CIFAR10 and ImageNet 64x64 experiments, we use the teacher model from [EDM](https://github.com/NVlabs/edm). The released checkpoint is a pickle file, so we need to extract the weights first. Run the official image sampling [code](https://github.com/NVlabs/edm/blob/main/generate.py) to save the model's state dict. 37 | 38 | We provide the extracted checkpoints for direct use under the same license as the original EDM checkpoint [Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License](http://creativecommons.org/licenses/by-nc-sa/4.0/): 39 | - [EDM-CIFAR10](https://drive.google.com/file/d/1UT72TxuDcJ6F54fsBgDZDVYix1sS8vKd/view?usp=sharing) 40 | - [EDM-ImageNet64x64-EDM](https://drive.google.com/file/d/1sKFMEk48BHb7x7FJpPHsLxTGyIhgpCOm/view?usp=sharing) 41 | 42 | ### Pretrained PID Checkpoints 43 | Additionally, our distilled models for these datasets are available for direct evaluation: 44 | - [PID-CIFAR10](https://drive.google.com/file/d/1uhJnW-vbdheHIMX2NyoqW921-f9VTuYI/view?usp=drive_link) 45 | - [PID-ImageNet64x64](https://drive.google.com/file/d/1crecnZxE8BwHSp8YaEzV2jpAEH4MYUb4/view?usp=drive_link) 46 | 47 | Download the checkpoints and place them in the ./model_zoo directory. 48 | 49 | ## 2. Distillation with PID 50 | To start the distillation, use the bash scripts: 51 | ```bash 52 | bash ./scripts/distill_pid_diffusion.sh 53 | ``` 54 | 55 | We use Open MPI to launch our code. Before running the experiment, configure the following in the bash file: 56 | 57 | > a. Set the environment variable `OPENAI_LOGDIR` to specify where the experiment data will be stored (e.g., `../experiment/EXP_NAME`, where `EXP_NAME` is the experiment name). 58 | > 59 | > b. Specify the number of GPUs to use (e.g., `-np 8` to use 8 GPUs). 60 | > 61 | > c. Set the total batch size across all GPUs (e.g., `--global_batch_size 512`, which will result in a batch size of `512/8=64` per GPU). 62 | 63 | ## 3. Image Sampling for EDM and PID model 64 | Use the bash script `./scripts/image_sampling.sh` to sample images from the pre-trained teacher model or the distilled model. The distilled PID model can be downloaded [here](https://drive.google.com/drive/folders/1rOmGWPyfhaVr6nfbVzJ8Xruk2ePWu1XE?usp=sharing). We provide the distilled one step model for both CIFAR and ImageNet64. 65 | 66 | Overview 67 | 68 | ## 4. FID Evaluation 69 | To evaluate FID scores, use the provided bash script `./scripts/fid_eval.sh`, which will evaluate all checkpoints in the `EXP_PATH` folder. Download the reference statistics for the teacher model from [EDM](https://nvlabs-fi-cdn.nvidia.com/edm/fid-refs/) and place them in `./model_zoo/stats/cifar10-32x32.npz` and `./model_zoo/stats/imagenet-64x64.npz`. Run the following to download the reference statistics: 70 | ```bash 71 | mkdir ./model_zoo/stats 72 | wget https://nvlabs-fi-cdn.nvidia.com/edm/fid-refs/cifar10-32x32.npz -P ./model_zoo/stats 73 | wget https://nvlabs-fi-cdn.nvidia.com/edm/fid-refs/imagenet-64x64.npz -o ./model_zoo/stats/imagenet-64x64.npz 74 | ``` 75 | 76 | To assess our pretrained [CIFAR10 model](https://drive.google.com/file/d/1uhJnW-vbdheHIMX2NyoqW921-f9VTuYI/view?usp=sharing), place it in `model_zoo/pid_cifar/pid_cifar.pt`, then execute the following for evaluation: 77 | ```bash 78 | EXP_PATH="./model_zoo/pid_cifar" 79 | 80 | mpirun -np 1 python ./scripts/fid_evaluation.py \ 81 | --training_mode one_shot_pinn_edm_edm_one_shot \ 82 | --fid_dataset cifar10 \ 83 | --exp_dir $EXP_PATH\ 84 | --batch_size 125 \ 85 | --sigma_max 80 \ 86 | --sigma_min 0.002 \ 87 | --s_churn 0 \ 88 | --steps 35 \ 89 | --sampler oneshot \ 90 | --attention_resolutions "2" \ 91 | --class_cond False \ 92 | --dropout 0.0 \ 93 | --image_size 32 \ 94 | --num_channels 128 \ 95 | --num_res_blocks 4 \ 96 | --num_samples 50000 \ 97 | --resblock_updown True \ 98 | --use_fp16 False \ 99 | --use_scale_shift_norm True \ 100 | --weight_schedule uniform \ 101 | --seed 0 102 | ``` 103 | 104 | # Citation 105 | ``` 106 | @article{ 107 | tee2024physics, 108 | title={Physics Informed Distillation for Diffusion Models}, 109 | author={Joshua Tian Jin Tee and Kang Zhang and Hee Suk Yoon and Dhananjaya Nagaraja Gowda and Chanwoo Kim and Chang D. Yoo}, 110 | journal={Transactions on Machine Learning Research}, 111 | issn={2835-8856}, 112 | year={2024}, 113 | url={https://openreview.net/forum?id=rOvaUsF996}, 114 | note={} 115 | } 116 | ``` 117 | 118 | # Acknowledgments 119 | This repository is based on [openai/consistency_models](https://github.com/openai/consistency_models) and [EDM](https://github.com/NVlabs/edm). 120 | 121 | -------------------------------------------------------------------------------- /scripts/cm_train.py: -------------------------------------------------------------------------------- 1 | """ 2 | Train a diffusion model on images. 3 | """ 4 | 5 | import argparse 6 | 7 | from cm import dist_util, logger 8 | from cm.image_datasets import load_data 9 | from cm.resample import create_named_schedule_sampler 10 | from cm.script_util import ( 11 | model_and_diffusion_defaults, 12 | create_model_and_diffusion, 13 | create_one_shot_edmedm_model_and_diffusion, 14 | cm_train_defaults, 15 | args_to_dict, 16 | add_dict_to_argparser, 17 | create_ema_and_scales_fn, 18 | ) 19 | from cm.train_util import ODETrainLoop 20 | import torch.distributed as dist 21 | import copy 22 | import torch 23 | 24 | def main(): 25 | args = create_argparser().parse_args() 26 | 27 | dist_util.setup_dist() 28 | logger.configure() 29 | 30 | logger.log("creating model and diffusion...") 31 | ema_scale_fn = create_ema_and_scales_fn( 32 | target_ema_mode=args.target_ema_mode, 33 | start_ema=args.start_ema, 34 | scale_mode=args.scale_mode, 35 | start_scales=args.start_scales, 36 | end_scales=args.end_scales, 37 | total_steps=args.total_training_steps, 38 | distill_steps_per_iter=args.distill_steps_per_iter, 39 | ) 40 | 41 | model_and_diffusion_kwargs = args_to_dict(args, model_and_diffusion_defaults().keys()) 42 | if args.training_mode == "progdist": 43 | distillation = False 44 | model_and_diffusion_kwargs["distillation"] = distillation 45 | model, diffusion = create_model_and_diffusion(**model_and_diffusion_kwargs) 46 | 47 | elif "consistency" in args.training_mode: 48 | distillation = True 49 | model_and_diffusion_kwargs["distillation"] = distillation 50 | model, diffusion = create_model_and_diffusion(**model_and_diffusion_kwargs) 51 | 52 | elif args.training_mode == "one_shot_pinn_edm_edm": 53 | student_model_and_diffusion_kwargs = copy.deepcopy(model_and_diffusion_kwargs) 54 | 55 | 56 | student_model_and_diffusion_kwargs["random_init"] = False 57 | 58 | model, diffusion = create_one_shot_edmedm_model_and_diffusion(**student_model_and_diffusion_kwargs) 59 | 60 | else: 61 | raise ValueError(f"unknown training mode {args.training_mode}") 62 | 63 | 64 | model.to(dist_util.dev()) 65 | model.train() 66 | 67 | if args.use_fp16: 68 | model.convert_to_fp16() 69 | 70 | # A distribution over timesteps in the diffusion process, intended to reduce 71 | # variance of the objective. 72 | schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion) 73 | 74 | if args.batch_size == -1: 75 | batch_size = args.global_batch_size // dist.get_world_size() 76 | if args.global_batch_size % dist.get_world_size() != 0: 77 | logger.log( 78 | f"warning, using smaller global_batch_size of {dist.get_world_size()*batch_size} instead of {args.global_batch_size}" 79 | ) 80 | else: 81 | batch_size = args.batch_size 82 | # batch_size = 2048 which is the batch_size for each gpu 83 | 84 | data = None 85 | 86 | if len(args.teacher_model_path) > 0: # path to the teacher score model. 87 | logger.log(f"loading the teacher model from {args.teacher_model_path}") 88 | teacher_model_and_diffusion_kwargs = copy.deepcopy(model_and_diffusion_kwargs) 89 | teacher_model_and_diffusion_kwargs["dropout"] = args.teacher_dropout # 0.1 90 | teacher_model_and_diffusion_kwargs["distillation"] = False 91 | teacher_model_and_diffusion_kwargs["random_init"] = False 92 | 93 | if args.training_mode == "one_shot_pinn_edm_edm": 94 | teacher_model_and_diffusion_kwargs["teacher_precond"] = True 95 | teacher_model, teacher_diffusion = create_one_shot_edmedm_model_and_diffusion(**teacher_model_and_diffusion_kwargs) 96 | load_weights = dist_util.load_state_dict(args.teacher_model_path, map_location="cpu") 97 | new_state_dict = {} 98 | for k, v in load_weights.items(): 99 | if "map_augment" in k: 100 | continue 101 | new_key = k.replace("model.", "") 102 | new_state_dict[new_key] = v 103 | teacher_model.load_state_dict(new_state_dict) 104 | 105 | model.load_state_dict(new_state_dict) 106 | 107 | else: 108 | teacher_model, teacher_diffusion = create_model_and_diffusion( 109 | **teacher_model_and_diffusion_kwargs, 110 | ) 111 | teacher_model.load_state_dict( 112 | dist_util.load_state_dict(args.teacher_model_path, map_location="cpu"), 113 | ) 114 | 115 | teacher_model.to(dist_util.dev()) 116 | teacher_model.eval() 117 | 118 | 119 | if args.use_fp16: 120 | teacher_model.convert_to_fp16() 121 | 122 | else: 123 | teacher_model = None 124 | teacher_diffusion = None 125 | 126 | logger.log("training...") 127 | ODETrainLoop( 128 | model=model, 129 | teacher_model=teacher_model, 130 | teacher_diffusion=teacher_diffusion, 131 | training_mode=args.training_mode, 132 | ema_scale_fn=ema_scale_fn, 133 | total_training_steps=args.total_training_steps, # 600000 134 | diffusion=diffusion, 135 | data=data, 136 | batch_size=batch_size, 137 | microbatch=args.microbatch, # -1 138 | lr=args.lr, # 1e-4 139 | ema_rate=args.ema_rate, # 0.999,0.9999,0.9999432189950708 140 | log_interval=args.log_interval, # 10 141 | save_interval=args.save_interval, # 10k 142 | resume_checkpoint=args.resume_checkpoint, 143 | use_fp16=args.use_fp16, # True 144 | fp16_scale_growth=args.fp16_scale_growth, # 1e-3 145 | schedule_sampler=schedule_sampler, 146 | weight_decay=args.weight_decay, # 0 147 | lr_anneal_steps=args.lr_anneal_steps, # 0 148 | methodology=args.methodology, 149 | optimizer=args.optimizer, 150 | opt_eps=args.opt_eps, 151 | eval_interval=args.eval_interval, 152 | ).run_loop() 153 | 154 | 155 | def create_argparser(): 156 | defaults = dict( 157 | data_dir="../cifar10/train", 158 | schedule_sampler="uniform", 159 | lr=1e-4, 160 | weight_decay=0.0, 161 | lr_anneal_steps=0, 162 | global_batch_size=2048, 163 | batch_size=-1, 164 | microbatch=-1, # -1 disables microbatches 165 | ema_rate="0.9999", # comma-separated list of EMA values 166 | log_interval=10, 167 | save_interval=10000, 168 | resume_checkpoint="", 169 | use_fp16=False, 170 | fp16_scale_growth=1e-3, 171 | methodology="Euler", 172 | use_target_model=False, 173 | optimizer='radam', 174 | opt_eps=1e-8, 175 | eval_interval=10000, 176 | random_init_stu=False, 177 | ) 178 | defaults.update(model_and_diffusion_defaults()) 179 | defaults.update(cm_train_defaults()) 180 | parser = argparse.ArgumentParser() 181 | add_dict_to_argparser(parser, defaults) 182 | return parser 183 | 184 | 185 | if __name__ == "__main__": 186 | main() 187 | -------------------------------------------------------------------------------- /cm/resample.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | 3 | import numpy as np 4 | import torch as th 5 | from scipy.stats import norm 6 | import torch.distributed as dist 7 | 8 | 9 | def create_named_schedule_sampler(name, diffusion): 10 | """ 11 | Create a ScheduleSampler from a library of pre-defined samplers. 12 | 13 | :param name: the name of the sampler. 14 | :param diffusion: the diffusion object to sample for. 15 | """ 16 | if name == "uniform": 17 | return UniformSampler(diffusion) 18 | elif name == "loss-second-moment": 19 | return LossSecondMomentResampler(diffusion) 20 | elif name == "lognormal": 21 | return LogNormalSampler() 22 | else: 23 | raise NotImplementedError(f"unknown schedule sampler: {name}") 24 | 25 | 26 | class ScheduleSampler(ABC): 27 | """ 28 | A distribution over timesteps in the diffusion process, intended to reduce 29 | variance of the objective. 30 | 31 | By default, samplers perform unbiased importance sampling, in which the 32 | objective's mean is unchanged. 33 | However, subclasses may override sample() to change how the resampled 34 | terms are reweighted, allowing for actual changes in the objective. 35 | """ 36 | 37 | @abstractmethod 38 | def weights(self): 39 | """ 40 | Get a numpy array of weights, one per diffusion step. 41 | 42 | The weights needn't be normalized, but must be positive. 43 | """ 44 | 45 | def sample(self, batch_size, device): 46 | """ 47 | Importance-sample timesteps for a batch. 48 | 49 | :param batch_size: the number of timesteps. 50 | :param device: the torch device to save to. 51 | :return: a tuple (timesteps, weights): 52 | - timesteps: a tensor of timestep indices. 53 | - weights: a tensor of weights to scale the resulting losses. 54 | """ 55 | w = self.weights() 56 | p = w / np.sum(w) 57 | indices_np = np.random.choice(len(p), size=(batch_size,), p=p) 58 | indices = th.from_numpy(indices_np).long().to(device) 59 | weights_np = 1 / (len(p) * p[indices_np]) 60 | weights = th.from_numpy(weights_np).float().to(device) 61 | return indices, weights 62 | 63 | 64 | class UniformSampler(ScheduleSampler): 65 | def __init__(self, diffusion): 66 | self.diffusion = diffusion 67 | self._weights = np.ones([diffusion.num_timesteps]) 68 | 69 | def weights(self): 70 | return self._weights 71 | 72 | 73 | class LossAwareSampler(ScheduleSampler): 74 | def update_with_local_losses(self, local_ts, local_losses): 75 | """ 76 | Update the reweighting using losses from a model. 77 | 78 | Call this method from each rank with a batch of timesteps and the 79 | corresponding losses for each of those timesteps. 80 | This method will perform synchronization to make sure all of the ranks 81 | maintain the exact same reweighting. 82 | 83 | :param local_ts: an integer Tensor of timesteps. 84 | :param local_losses: a 1D Tensor of losses. 85 | """ 86 | batch_sizes = [ 87 | th.tensor([0], dtype=th.int32, device=local_ts.device) 88 | for _ in range(dist.get_world_size()) 89 | ] 90 | dist.all_gather( 91 | batch_sizes, 92 | th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device), 93 | ) 94 | 95 | # Pad all_gather batches to be the maximum batch size. 96 | batch_sizes = [x.item() for x in batch_sizes] 97 | max_bs = max(batch_sizes) 98 | 99 | timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes] 100 | loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes] 101 | dist.all_gather(timestep_batches, local_ts) 102 | dist.all_gather(loss_batches, local_losses) 103 | timesteps = [ 104 | x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs] 105 | ] 106 | losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]] 107 | self.update_with_all_losses(timesteps, losses) 108 | 109 | @abstractmethod 110 | def update_with_all_losses(self, ts, losses): 111 | """ 112 | Update the reweighting using losses from a model. 113 | 114 | Sub-classes should override this method to update the reweighting 115 | using losses from the model. 116 | 117 | This method directly updates the reweighting without synchronizing 118 | between workers. It is called by update_with_local_losses from all 119 | ranks with identical arguments. Thus, it should have deterministic 120 | behavior to maintain state across workers. 121 | 122 | :param ts: a list of int timesteps. 123 | :param losses: a list of float losses, one per timestep. 124 | """ 125 | 126 | 127 | class LossSecondMomentResampler(LossAwareSampler): 128 | def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001): 129 | self.diffusion = diffusion 130 | self.history_per_term = history_per_term 131 | self.uniform_prob = uniform_prob 132 | self._loss_history = np.zeros( 133 | [diffusion.num_timesteps, history_per_term], dtype=np.float64 134 | ) 135 | self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int) 136 | 137 | def weights(self): 138 | if not self._warmed_up(): 139 | return np.ones([self.diffusion.num_timesteps], dtype=np.float64) 140 | weights = np.sqrt(np.mean(self._loss_history**2, axis=-1)) 141 | weights /= np.sum(weights) 142 | weights *= 1 - self.uniform_prob 143 | weights += self.uniform_prob / len(weights) 144 | return weights 145 | 146 | def update_with_all_losses(self, ts, losses): 147 | for t, loss in zip(ts, losses): 148 | if self._loss_counts[t] == self.history_per_term: 149 | # Shift out the oldest loss term. 150 | self._loss_history[t, :-1] = self._loss_history[t, 1:] 151 | self._loss_history[t, -1] = loss 152 | else: 153 | self._loss_history[t, self._loss_counts[t]] = loss 154 | self._loss_counts[t] += 1 155 | 156 | def _warmed_up(self): 157 | return (self._loss_counts == self.history_per_term).all() 158 | 159 | 160 | class LogNormalSampler: 161 | def __init__(self, p_mean=-1.2, p_std=1.2, even=False): 162 | self.p_mean = p_mean 163 | self.p_std = p_std 164 | self.even = even 165 | if self.even: 166 | self.inv_cdf = lambda x: norm.ppf(x, loc=p_mean, scale=p_std) 167 | self.rank, self.size = dist.get_rank(), dist.get_world_size() 168 | 169 | def sample(self, bs, device): 170 | if self.even: 171 | # buckets = [1/G] 172 | start_i, end_i = self.rank * bs, (self.rank + 1) * bs 173 | global_batch_size = self.size * bs 174 | locs = (th.arange(start_i, end_i) + th.rand(bs)) / global_batch_size 175 | log_sigmas = th.tensor(self.inv_cdf(locs), dtype=th.float32, device=device) 176 | else: 177 | log_sigmas = self.p_mean + self.p_std * th.randn(bs, device=device) 178 | sigmas = th.exp(log_sigmas) 179 | weights = th.ones_like(sigmas) 180 | return sigmas, weights 181 | -------------------------------------------------------------------------------- /cm/fp16_util.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers to train with 16-bit precision. 3 | """ 4 | 5 | import numpy as np 6 | import torch as th 7 | import torch.nn as nn 8 | from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors 9 | 10 | from . import logger 11 | 12 | INITIAL_LOG_LOSS_SCALE = 20.0 13 | # INITIAL_LOG_LOSS_SCALE = 12.0 14 | 15 | 16 | def convert_module_to_f16(l): 17 | """ 18 | Convert primitive modules to float16. 19 | """ 20 | if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): 21 | l.weight.data = l.weight.data.half() 22 | if l.bias is not None: 23 | l.bias.data = l.bias.data.half() 24 | 25 | 26 | def convert_module_to_f32(l): 27 | """ 28 | Convert primitive modules to float32, undoing convert_module_to_f16(). 29 | """ 30 | if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): 31 | l.weight.data = l.weight.data.float() 32 | if l.bias is not None: 33 | l.bias.data = l.bias.data.float() 34 | 35 | 36 | def make_master_params(param_groups_and_shapes): 37 | """ 38 | Copy model parameters into a (differently-shaped) list of full-precision 39 | parameters. 40 | """ 41 | master_params = [] 42 | for param_group, shape in param_groups_and_shapes: 43 | master_param = nn.Parameter( 44 | _flatten_dense_tensors( 45 | [param.detach().float() for (_, param) in param_group] 46 | ).view(shape) 47 | ) 48 | master_param.requires_grad = True 49 | master_params.append(master_param) 50 | return master_params 51 | 52 | 53 | def model_grads_to_master_grads(param_groups_and_shapes, master_params): 54 | """ 55 | Copy the gradients from the model parameters into the master parameters 56 | from make_master_params(). 57 | """ 58 | for master_param, (param_group, shape) in zip( 59 | master_params, param_groups_and_shapes 60 | ): 61 | master_param.grad = _flatten_dense_tensors( 62 | [param_grad_or_zeros(param) for (_, param) in param_group] 63 | ).view(shape) 64 | 65 | 66 | def master_params_to_model_params(param_groups_and_shapes, master_params): 67 | """ 68 | Copy the master parameter data back into the model parameters. 69 | """ 70 | # Without copying to a list, if a generator is passed, this will 71 | # silently not copy any parameters. 72 | for master_param, (param_group, _) in zip(master_params, param_groups_and_shapes): 73 | for (_, param), unflat_master_param in zip( 74 | param_group, unflatten_master_params(param_group, master_param.view(-1)) 75 | ): 76 | param.detach().copy_(unflat_master_param) 77 | 78 | 79 | def unflatten_master_params(param_group, master_param): 80 | return _unflatten_dense_tensors(master_param, [param for (_, param) in param_group]) 81 | 82 | 83 | def get_param_groups_and_shapes(named_model_params): 84 | named_model_params = list(named_model_params) 85 | scalar_vector_named_params = ( 86 | [(n, p) for (n, p) in named_model_params if p.ndim <= 1], 87 | (-1), 88 | ) 89 | matrix_named_params = ( 90 | [(n, p) for (n, p) in named_model_params if p.ndim > 1], 91 | (1, -1), 92 | ) 93 | return [scalar_vector_named_params, matrix_named_params] 94 | 95 | 96 | def master_params_to_state_dict( 97 | model, param_groups_and_shapes, master_params, use_fp16 98 | ): 99 | if use_fp16: 100 | state_dict = model.state_dict() 101 | for master_param, (param_group, _) in zip( 102 | master_params, param_groups_and_shapes 103 | ): 104 | for (name, _), unflat_master_param in zip( 105 | param_group, unflatten_master_params(param_group, master_param.view(-1)) 106 | ): 107 | assert name in state_dict 108 | state_dict[name] = unflat_master_param 109 | else: 110 | state_dict = model.state_dict() 111 | for i, (name, _value) in enumerate(model.named_parameters()): 112 | assert name in state_dict 113 | state_dict[name] = master_params[i] 114 | return state_dict 115 | 116 | 117 | def state_dict_to_master_params(model, state_dict, use_fp16): 118 | if use_fp16: 119 | named_model_params = [ 120 | (name, state_dict[name]) for name, _ in model.named_parameters() 121 | ] 122 | param_groups_and_shapes = get_param_groups_and_shapes(named_model_params) 123 | master_params = make_master_params(param_groups_and_shapes) 124 | else: 125 | master_params = [state_dict[name] for name, _ in model.named_parameters()] 126 | return master_params 127 | 128 | 129 | def zero_master_grads(master_params): 130 | for param in master_params: 131 | param.grad = None 132 | 133 | 134 | def zero_grad(model_params): 135 | for param in model_params: 136 | # Taken from https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.add_param_group 137 | if param.grad is not None: 138 | param.grad.detach_() 139 | param.grad.zero_() 140 | 141 | 142 | def param_grad_or_zeros(param): 143 | if param.grad is not None: 144 | return param.grad.data.detach() 145 | else: 146 | return th.zeros_like(param) 147 | 148 | 149 | class MixedPrecisionTrainer: 150 | def __init__( 151 | self, 152 | *, 153 | model, 154 | use_fp16=False, 155 | fp16_scale_growth=1e-3, 156 | initial_lg_loss_scale=INITIAL_LOG_LOSS_SCALE, 157 | ): 158 | self.model = model 159 | self.use_fp16 = use_fp16 160 | self.fp16_scale_growth = fp16_scale_growth 161 | 162 | self.model_params = list(self.model.parameters()) 163 | 164 | self.master_params = self.model_params 165 | self.param_groups_and_shapes = None 166 | self.lg_loss_scale = initial_lg_loss_scale 167 | 168 | if self.use_fp16: 169 | self.param_groups_and_shapes = get_param_groups_and_shapes( 170 | self.model.named_parameters() 171 | ) 172 | self.master_params = make_master_params(self.param_groups_and_shapes) 173 | self.model.convert_to_fp16() 174 | 175 | def zero_grad(self): 176 | zero_grad(self.model_params) 177 | 178 | def backward(self, loss: th.Tensor): 179 | if self.use_fp16: 180 | loss_scale = 2**self.lg_loss_scale 181 | (loss * loss_scale).backward() 182 | else: 183 | loss.backward() 184 | 185 | def optimize(self, opt: th.optim.Optimizer): 186 | if self.use_fp16: 187 | return self._optimize_fp16(opt) 188 | else: 189 | return self._optimize_normal(opt) 190 | 191 | def _optimize_fp16(self, opt: th.optim.Optimizer): 192 | logger.logkv_mean("lg_loss_scale", self.lg_loss_scale) 193 | model_grads_to_master_grads(self.param_groups_and_shapes, self.master_params) 194 | grad_norm, param_norm = self._compute_norms(grad_scale=2**self.lg_loss_scale) 195 | if check_overflow(grad_norm): 196 | self.lg_loss_scale -= 1 197 | logger.log(f"Found NaN, decreased lg_loss_scale to {self.lg_loss_scale}") 198 | zero_master_grads(self.master_params) 199 | return False 200 | 201 | logger.logkv_mean("grad_norm", grad_norm) 202 | logger.logkv_mean("param_norm", param_norm) 203 | 204 | for p in self.master_params: 205 | p.grad.mul_(1.0 / (2**self.lg_loss_scale)) 206 | # th.nn.utils.clip_grad_norm_(self.master_params, 1.) 207 | opt.step() 208 | zero_master_grads(self.master_params) 209 | master_params_to_model_params(self.param_groups_and_shapes, self.master_params) 210 | self.lg_loss_scale += self.fp16_scale_growth 211 | return True 212 | 213 | def _optimize_normal(self, opt: th.optim.Optimizer): 214 | grad_norm, param_norm = self._compute_norms() 215 | logger.logkv_mean("grad_norm", grad_norm) 216 | logger.logkv_mean("param_norm", param_norm) 217 | # th.nn.utils.clip_grad_norm_(self.master_params, 1.) 218 | opt.step() 219 | return True 220 | 221 | def _compute_norms(self, grad_scale=1.0): 222 | grad_norm = 0.0 223 | param_norm = 0.0 224 | for p in self.master_params: 225 | with th.no_grad(): 226 | param_norm += th.norm(p, p=2, dtype=th.float32).item() ** 2 227 | if p.grad is not None: 228 | grad_norm += th.norm(p.grad, p=2, dtype=th.float32).item() ** 2 229 | return np.sqrt(grad_norm) / grad_scale, np.sqrt(param_norm) 230 | 231 | def master_params_to_state_dict(self, master_params): 232 | return master_params_to_state_dict( 233 | self.model, self.param_groups_and_shapes, master_params, self.use_fp16 234 | ) 235 | 236 | def state_dict_to_master_params(self, state_dict): 237 | return state_dict_to_master_params(self.model, state_dict, self.use_fp16) 238 | 239 | def check_overflow(value): 240 | return (value == float("inf")) or (value == -float("inf")) or (value != value) 241 | -------------------------------------------------------------------------------- /scripts/fid_evaluation.py: -------------------------------------------------------------------------------- 1 | """ 2 | Generate a large batch of image samples from a model and save them as a large 3 | numpy array. This can be used to produce samples for FID evaluation. 4 | """ 5 | 6 | import argparse 7 | import os 8 | 9 | import numpy as np 10 | import torch as th 11 | import torch.distributed as dist 12 | 13 | from cm import dist_util, logger 14 | from cm.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 | create_one_shot_edmedm_model_and_diffusion, 21 | ) 22 | from cm.random_util import get_generator 23 | from cm.karras_diffusion import karras_sample 24 | import torchvision 25 | import PIL 26 | from cleanfid.features import build_feature_extractor, get_reference_statistics 27 | import scipy.linalg 28 | import glob 29 | import csv 30 | from pathlib import Path 31 | import pandas 32 | 33 | def calculate_fid_from_inception_stats(mu, sigma, mu_ref, sigma_ref): 34 | m = np.square(mu - mu_ref).sum() 35 | s, _ = scipy.linalg.sqrtm(np.dot(sigma, sigma_ref), disp=False) 36 | fid = m + np.trace(sigma + sigma_ref - s * 2) 37 | return float(np.real(fid)) 38 | 39 | def main(args, model=None, diffusion=None): 40 | 41 | if model==None or diffusion==None: 42 | 43 | if "one_shot_pinn_edm_edm_teacher" == args.training_mode: 44 | model, diffusion = create_one_shot_edmedm_model_and_diffusion(teacher_precond=True, 45 | **args_to_dict(args, model_and_diffusion_defaults().keys()), 46 | ) 47 | load_weights = dist_util.load_state_dict(args.model_path, map_location="cpu") 48 | new_state_dict = {} 49 | for k, v in load_weights.items(): 50 | if "map_augment" in k: 51 | continue 52 | new_key = k.replace("model.", "") 53 | new_state_dict[new_key] = v 54 | 55 | model.load_state_dict(new_state_dict) 56 | elif "one_shot_pinn_edm_edm_one_shot" == args.training_mode: 57 | model, diffusion = create_one_shot_edmedm_model_and_diffusion(teacher_precond=False, 58 | **args_to_dict(args, model_and_diffusion_defaults().keys()), 59 | ) 60 | load_weights = dist_util.load_state_dict(args.model_path, map_location="cpu") 61 | new_state_dict = {} 62 | for k, v in load_weights.items(): 63 | if "map_augment" in k: 64 | continue 65 | new_key = k.replace("model.", "") 66 | new_state_dict[new_key] = v 67 | 68 | model.load_state_dict(new_state_dict) 69 | else: 70 | raise ValueError(f"training mode {args.training_mode} not supported") 71 | 72 | model.to(dist_util.dev()) 73 | if args.use_fp16: 74 | model.convert_to_fp16() 75 | model.eval() 76 | 77 | logger.log("sampling...") 78 | if args.sampler == "multistep": 79 | assert len(args.ts) > 0 80 | ts = tuple(int(x) for x in args.ts.split(",")) 81 | else: 82 | ts = None 83 | 84 | 85 | mode='legacy_tensorflow' 86 | fid_evaluation_feat_model = build_feature_extractor(mode, dist_util.dev(), use_dataparallel=False) 87 | 88 | if dist.get_rank() == 0: 89 | dataset_name=args.fid_dataset 90 | if dataset_name == 'cifar10': 91 | fpath = './model_zoo/stats/cifar10-32x32.npz' 92 | elif dataset_name == 'imagenet': 93 | fpath = './model_zoo/stats/imagenet-64x64.npz' 94 | else: 95 | raise ValueError(f"fid evaluation error: not support dataset {dataset_name}.") 96 | stats = np.load(fpath) 97 | ref_mu, ref_sigma = stats["mu"], stats["sigma"] 98 | 99 | 100 | all_images = [] 101 | all_labels = [] 102 | generator = get_generator(args.generator, args.num_samples, args.seed) 103 | 104 | feature_dim = 2048 105 | mu = th.zeros([feature_dim], dtype=th.float64, device=dist_util.dev()) 106 | sigma = th.zeros([feature_dim, feature_dim], dtype=th.float64, device=dist_util.dev()) 107 | l_feats = [] 108 | 109 | generated_img = False 110 | while len(all_images) * args.batch_size < args.num_samples: 111 | model_kwargs = {} 112 | if args.class_cond: 113 | classes = th.randint( 114 | low=0, high=NUM_CLASSES, size=(args.batch_size,), device=dist_util.dev() 115 | ) 116 | model_kwargs["y"] = classes 117 | 118 | sample_ori = karras_sample( 119 | diffusion, 120 | model, 121 | (args.batch_size, 3, args.image_size, args.image_size), 122 | steps=args.steps, 123 | model_kwargs=model_kwargs, 124 | device=dist_util.dev(), 125 | clip_denoised=args.clip_denoised, 126 | sampler=args.sampler, 127 | sigma_min=args.sigma_min, 128 | sigma_max=args.sigma_max, 129 | s_churn=args.s_churn, 130 | s_tmin=args.s_tmin, 131 | s_tmax=args.s_tmax, 132 | s_noise=args.s_noise, 133 | generator=generator, 134 | ts=ts, 135 | ) 136 | sample = ((sample_ori + 1) * 127.5).clamp(0, 255).to(th.uint8) 137 | sample = sample.permute(0, 2, 3, 1) 138 | sample = sample.contiguous() 139 | 140 | gathered_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())] 141 | dist.all_gather(gathered_samples, sample) # gather not supported with NCCL 142 | all_images.extend([sample.cpu().numpy() for sample in gathered_samples]) 143 | if args.class_cond: 144 | gathered_labels = [ 145 | th.zeros_like(classes) for _ in range(dist.get_world_size()) 146 | ] 147 | dist.all_gather(gathered_labels, classes) 148 | all_labels.extend([labels.cpu().numpy() for labels in gathered_labels]) 149 | logger.log(f"created {len(all_images) * args.batch_size} samples") 150 | 151 | with th.no_grad(): 152 | feat = fid_evaluation_feat_model(((sample_ori + 1.) * 0.5).mul(255).clip(0, 255).to(dist_util.dev())).to(th.float64) 153 | # l_feats.append(feat.cpu()) 154 | mu += feat.sum(0) 155 | sigma += feat.T @ feat 156 | 157 | if not generated_img and dist.get_rank() == 0: 158 | generated_img = True 159 | torchvision.utils.save_image((sample_ori+1.)/2., os.path.join(logger.get_dir(), f"{args.model_name}_samples.png"), nrow = 10) 160 | 161 | 162 | mu = mu.unsqueeze(0) 163 | sigma = sigma.unsqueeze(0) 164 | 165 | all_mu = [th.zeros_like(mu) for _ in range(dist.get_world_size())] 166 | dist.all_gather(all_mu, mu) # gather not supported with NCCL 167 | all_mu = th.cat(all_mu, axis=0) 168 | 169 | all_sigma = [th.zeros_like(sigma) for _ in range(dist.get_world_size())] 170 | dist.all_gather(all_sigma, sigma) # gather not supported with NCCL 171 | all_sigma = th.cat(all_sigma, axis=0) 172 | 173 | mu = all_mu.sum(0) 174 | sigma = all_sigma.sum(0) 175 | 176 | num_images = args.num_samples 177 | mu /= num_images 178 | sigma -= mu.ger(mu) * num_images 179 | sigma /= num_images - 1 180 | 181 | mu = mu.cpu().numpy() 182 | sigma = sigma.cpu().numpy() 183 | 184 | 185 | arr = np.concatenate(all_images, axis=0) 186 | arr = arr[: args.num_samples] 187 | if args.class_cond: 188 | label_arr = np.concatenate(all_labels, axis=0) 189 | label_arr = label_arr[: args.num_samples] 190 | fid_score = 0. 191 | if dist.get_rank() == 0: 192 | 193 | shape_str = "x".join([str(x) for x in arr.shape]) 194 | out_path = os.path.join(logger.get_dir(), f"{args.model_name}_samples_{shape_str}.npz") 195 | logger.log(f"saving to {out_path}") 196 | 197 | if args.class_cond: 198 | np.savez(out_path, arr, label_arr) 199 | else: 200 | np.savez(out_path, arr) 201 | 202 | 203 | fid_score = calculate_fid_from_inception_stats( 204 | mu, sigma, 205 | ref_mu, 206 | ref_sigma 207 | ) 208 | 209 | dist.barrier() 210 | logger.log("sampling complete") 211 | return fid_score 212 | 213 | 214 | def create_argparser(): 215 | defaults = dict( 216 | training_mode="edm", 217 | generator="determ", 218 | clip_denoised=True, 219 | num_samples=10000, 220 | batch_size=16, 221 | sampler="heun", 222 | s_churn=0.0, 223 | s_tmin=0.0, 224 | s_tmax=float("inf"), 225 | s_noise=1.0, 226 | steps=40, 227 | model_path="", 228 | seed=42, 229 | ts="", 230 | fid_dataset="cifar10", 231 | ) 232 | defaults.update(model_and_diffusion_defaults()) 233 | parser = argparse.ArgumentParser() 234 | parser.add_argument('--exp_dir', type=str) 235 | add_dict_to_argparser(parser, defaults) 236 | return parser 237 | 238 | 239 | if __name__ == "__main__": 240 | args = create_argparser().parse_args() 241 | exp_dir = Path(args.exp_dir) 242 | os.environ["OPENAI_LOGDIR"] = str(exp_dir / "FID") 243 | 244 | dist_util.setup_dist() 245 | logger.configure() 246 | assert args.num_samples % (args.batch_size * dist.get_world_size()) == 0 247 | 248 | result_file = exp_dir / f"fid_results_{args.seed}.csv" 249 | if result_file.exists(): 250 | logger.log("exist results, check result...") 251 | df = pandas.read_csv(str(result_file)) 252 | tested_model = list(df["model_name"]) 253 | logger.log(f"Tested model: {tested_model}") 254 | else: 255 | tested_model = [] 256 | if dist.get_rank() == 0: 257 | with open(str(result_file), mode='w') as f: 258 | fid_writer = csv.writer(f, delimiter=',') 259 | fid_writer.writerow(['model_name', 'FID']) 260 | 261 | model_list_ = exp_dir.glob("*.pt") 262 | model_list = [] 263 | 264 | for model_dir in model_list_: 265 | model_name = str(model_dir.stem) 266 | model_dir = str(model_dir) 267 | model_list.append( 268 | { 269 | "model_name": model_name, 270 | "path": model_dir, 271 | } 272 | ) 273 | 274 | for model_test in model_list: 275 | args.model_path = model_test["path"] 276 | args.model_name = model_test["model_name"] 277 | logger.log(f"\nmodel: {args.model_name}") 278 | 279 | if args.model_name in tested_model: 280 | continue 281 | 282 | fid_score = main(args) 283 | 284 | logger.log(f"\nmodel: {args.model_name} \nFID {fid_score:.4f}\n") 285 | 286 | if dist.get_rank() == 0: 287 | with open(str(result_file), mode='a') as f: 288 | fid_writer = csv.writer(f, delimiter=',') 289 | fid_writer.writerow([args.model_name, fid_score]) 290 | 291 | -------------------------------------------------------------------------------- /cm/script_util.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | from .karras_diffusion import KarrasDenoiser, OneShotDenoiser, EDMEDMDenoiser 4 | from .unet import UNetModel 5 | import numpy as np 6 | 7 | NUM_CLASSES = 1000 8 | 9 | 10 | def cm_train_defaults(): 11 | return dict( 12 | teacher_model_path="", 13 | teacher_dropout=0.1, 14 | training_mode="consistency_distillation", 15 | target_ema_mode="fixed", 16 | scale_mode="fixed", 17 | total_training_steps=600000, 18 | start_ema=0.0, 19 | start_scales=40, 20 | end_scales=40, 21 | distill_steps_per_iter=50000, 22 | loss_norm="lpips", 23 | ) 24 | 25 | 26 | def model_and_diffusion_defaults(): 27 | """ 28 | Defaults for image training. 29 | """ 30 | res = dict( 31 | sigma_min=0.002, 32 | sigma_max=80.0, 33 | image_size=64, 34 | num_channels=128, 35 | num_res_blocks=2, 36 | num_heads=4, 37 | num_heads_upsample=-1, 38 | num_head_channels=-1, 39 | attention_resolutions="32,16,8", 40 | channel_mult="", 41 | dropout=0.0, 42 | class_cond=False, 43 | use_checkpoint=False, 44 | use_scale_shift_norm=True, 45 | resblock_updown=False, 46 | use_fp16=False, 47 | use_new_attention_order=False, 48 | learn_sigma=False, 49 | weight_schedule="karras", 50 | loss_norm="lpips", 51 | continuous=False, 52 | ) 53 | return res 54 | 55 | 56 | def create_model_and_diffusion( 57 | image_size, 58 | class_cond, 59 | learn_sigma, 60 | num_channels, 61 | num_res_blocks, 62 | channel_mult, 63 | num_heads, 64 | num_head_channels, 65 | num_heads_upsample, 66 | attention_resolutions, 67 | dropout, 68 | use_checkpoint, 69 | use_scale_shift_norm, 70 | resblock_updown, 71 | use_fp16, 72 | use_new_attention_order, 73 | weight_schedule, 74 | sigma_min=0.002, 75 | sigma_max=80.0, 76 | distillation=False, 77 | loss_norm=None, 78 | continuous=None, 79 | ): 80 | model = create_model( 81 | image_size, 82 | num_channels, 83 | num_res_blocks, 84 | channel_mult=channel_mult, 85 | learn_sigma=learn_sigma, 86 | class_cond=class_cond, 87 | use_checkpoint=use_checkpoint, 88 | attention_resolutions=attention_resolutions, 89 | num_heads=num_heads, 90 | num_head_channels=num_head_channels, 91 | num_heads_upsample=num_heads_upsample, 92 | use_scale_shift_norm=use_scale_shift_norm, 93 | dropout=dropout, 94 | resblock_updown=resblock_updown, 95 | use_fp16=use_fp16, 96 | use_new_attention_order=use_new_attention_order, 97 | ) 98 | diffusion = KarrasDenoiser( 99 | sigma_data=0.5, 100 | sigma_max=sigma_max, 101 | sigma_min=sigma_min, 102 | distillation=distillation, 103 | weight_schedule=weight_schedule, 104 | ) 105 | return model, diffusion 106 | 107 | 108 | def create_one_shot_edmedm_model_and_diffusion( 109 | image_size, 110 | class_cond, 111 | learn_sigma, 112 | num_channels, 113 | num_res_blocks, 114 | channel_mult, 115 | num_heads, 116 | num_head_channels, 117 | num_heads_upsample, 118 | attention_resolutions, 119 | dropout, 120 | use_checkpoint, 121 | use_scale_shift_norm, 122 | resblock_updown, 123 | use_fp16, 124 | use_new_attention_order, 125 | weight_schedule, 126 | sigma_min=0.002, 127 | sigma_max=80.0, 128 | distillation=False, 129 | teacher_precond=False, 130 | loss_norm="l2", 131 | random_init=False, 132 | continuous=False, 133 | ): 134 | from cm.network import DhariwalUNet, SongUNet 135 | channel_mult__ = channel_mult 136 | 137 | if channel_mult__ == "": 138 | if image_size == 512: 139 | channel_mult = (0.5, 1, 1, 2, 2, 4, 4) 140 | elif image_size == 256: 141 | channel_mult = (1, 1, 2, 2, 4, 4) 142 | elif image_size == 128: 143 | channel_mult = (1, 1, 2, 3, 4) 144 | elif image_size == 64: 145 | channel_mult = (1, 2, 3, 4) 146 | elif image_size == 32: 147 | channel_mult = (2,2,2) 148 | else: 149 | raise ValueError(f"unsupported image size: {image_size}") 150 | else: 151 | channel_mult = tuple(int(ch_mult) for ch_mult in channel_mult__.split(",")) 152 | 153 | attention_ds = [] 154 | for res in attention_resolutions.split(","): 155 | attention_ds.append(image_size // int(res)) 156 | 157 | t_rescale = 1. 158 | if image_size == 32: 159 | model = SongUNet( 160 | image_size=image_size, 161 | in_channels=3, 162 | model_channels=num_channels, 163 | out_channels=(3 if not learn_sigma else 6), 164 | num_res_blocks=num_res_blocks, 165 | attention_resolutions=tuple(attention_ds), 166 | dropout=dropout, 167 | channel_mult=channel_mult, 168 | num_classes=(10 if class_cond else None), 169 | use_checkpoint=use_checkpoint, 170 | use_fp16=use_fp16, 171 | num_heads=num_heads, 172 | num_head_channels=num_head_channels, 173 | num_heads_upsample=num_heads_upsample, 174 | use_scale_shift_norm=use_scale_shift_norm, 175 | resblock_updown=resblock_updown, 176 | use_new_attention_order=use_new_attention_order, 177 | random_init=random_init, 178 | ) 179 | elif image_size == 64: 180 | model = DhariwalUNet( 181 | image_size=image_size, 182 | in_channels=3, 183 | model_channels=num_channels, 184 | out_channels=(3 if not learn_sigma else 6), 185 | num_res_blocks=num_res_blocks, 186 | attention_resolutions=tuple(attention_ds), 187 | dropout=dropout, 188 | channel_mult=channel_mult, 189 | num_classes=(NUM_CLASSES if class_cond else None), 190 | use_checkpoint=use_checkpoint, 191 | use_fp16=use_fp16, 192 | num_heads=num_heads, 193 | num_head_channels=num_head_channels, 194 | num_heads_upsample=num_heads_upsample, 195 | use_scale_shift_norm=use_scale_shift_norm, 196 | resblock_updown=resblock_updown, 197 | use_new_attention_order=use_new_attention_order, 198 | ) 199 | 200 | 201 | if teacher_precond: 202 | diffusion = EDMEDMDenoiser( 203 | sigma_data=0.5, 204 | sigma_max=sigma_max, 205 | sigma_min=sigma_min, 206 | distillation=distillation, 207 | weight_schedule=weight_schedule, 208 | t_rescale=t_rescale, 209 | ) 210 | else: 211 | diffusion = OneShotDenoiser( 212 | sigma_data=0.5, 213 | sigma_max=sigma_max, 214 | sigma_min=sigma_min, 215 | distillation=distillation, 216 | weight_schedule=weight_schedule, 217 | loss_norm=loss_norm, 218 | t_rescale=t_rescale, 219 | ) 220 | return model, diffusion 221 | 222 | def create_model( 223 | image_size, 224 | num_channels, 225 | num_res_blocks, 226 | channel_mult="", 227 | learn_sigma=False, 228 | class_cond=False, 229 | use_checkpoint=False, 230 | attention_resolutions="16", 231 | num_heads=1, 232 | num_head_channels=-1, 233 | num_heads_upsample=-1, 234 | use_scale_shift_norm=False, 235 | dropout=0, 236 | resblock_updown=False, 237 | use_fp16=False, 238 | use_new_attention_order=False, 239 | ): 240 | if channel_mult == "": 241 | if image_size == 512: 242 | channel_mult = (0.5, 1, 1, 2, 2, 4, 4) 243 | elif image_size == 256: 244 | channel_mult = (1, 1, 2, 2, 4, 4) 245 | elif image_size == 128: 246 | channel_mult = (1, 1, 2, 3, 4) 247 | elif image_size == 64: 248 | channel_mult = (1, 2, 3, 4) 249 | else: 250 | raise ValueError(f"unsupported image size: {image_size}") 251 | else: 252 | channel_mult = tuple(int(ch_mult) for ch_mult in channel_mult.split(",")) 253 | 254 | attention_ds = [] 255 | for res in attention_resolutions.split(","): 256 | attention_ds.append(image_size // int(res)) 257 | 258 | return UNetModel( 259 | image_size=image_size, 260 | in_channels=3, 261 | model_channels=num_channels, 262 | out_channels=(3 if not learn_sigma else 6), 263 | num_res_blocks=num_res_blocks, 264 | attention_resolutions=tuple(attention_ds), 265 | dropout=dropout, 266 | channel_mult=channel_mult, 267 | num_classes=(NUM_CLASSES if class_cond else None), 268 | use_checkpoint=use_checkpoint, 269 | use_fp16=use_fp16, 270 | num_heads=num_heads, 271 | num_head_channels=num_head_channels, 272 | num_heads_upsample=num_heads_upsample, 273 | use_scale_shift_norm=use_scale_shift_norm, 274 | resblock_updown=resblock_updown, 275 | use_new_attention_order=use_new_attention_order, 276 | ) 277 | 278 | 279 | def create_ema_and_scales_fn( 280 | target_ema_mode, 281 | start_ema, 282 | scale_mode, 283 | start_scales, 284 | end_scales, 285 | total_steps, 286 | distill_steps_per_iter, 287 | ): 288 | def ema_and_scales_fn(step): 289 | if target_ema_mode == "fixed" and scale_mode == "fixed": 290 | target_ema = start_ema 291 | scales = start_scales 292 | elif target_ema_mode == "fixed" and scale_mode == "progressive": 293 | target_ema = start_ema 294 | scales = np.ceil( 295 | np.sqrt( 296 | (step / total_steps) * ((end_scales + 1) ** 2 - start_scales**2) 297 | + start_scales**2 298 | ) 299 | - 1 300 | ).astype(np.int32) 301 | scales = np.maximum(scales, 1) 302 | scales = scales + 1 303 | 304 | elif target_ema_mode == "adaptive" and scale_mode == "progressive": 305 | scales = np.ceil( 306 | np.sqrt( 307 | (step / total_steps) * ((end_scales + 1) ** 2 - start_scales**2) 308 | + start_scales**2 309 | ) 310 | - 1 311 | ).astype(np.int32) 312 | scales = np.maximum(scales, 1) 313 | c = -np.log(start_ema) * start_scales 314 | target_ema = np.exp(-c / scales) 315 | scales = scales + 1 316 | elif target_ema_mode == "fixed" and scale_mode == "progdist": 317 | distill_stage = step // distill_steps_per_iter 318 | scales = start_scales // (2**distill_stage) 319 | scales = np.maximum(scales, 2) 320 | 321 | sub_stage = np.maximum( 322 | step - distill_steps_per_iter * (np.log2(start_scales) - 1), 323 | 0, 324 | ) 325 | sub_stage = sub_stage // (distill_steps_per_iter * 2) 326 | sub_scales = 2 // (2**sub_stage) 327 | sub_scales = np.maximum(sub_scales, 1) 328 | 329 | scales = np.where(scales == 2, sub_scales, scales) 330 | 331 | target_ema = 1.0 332 | else: 333 | raise NotImplementedError 334 | 335 | return float(target_ema), int(scales) 336 | 337 | return ema_and_scales_fn 338 | 339 | 340 | def add_dict_to_argparser(parser, default_dict): 341 | for k, v in default_dict.items(): 342 | v_type = type(v) 343 | if v is None: 344 | v_type = str 345 | elif isinstance(v, bool): 346 | v_type = str2bool 347 | parser.add_argument(f"--{k}", default=v, type=v_type) 348 | 349 | 350 | def args_to_dict(args, keys): 351 | return {k: getattr(args, k) for k in keys} 352 | 353 | 354 | def str2bool(v): 355 | """ 356 | https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse 357 | """ 358 | if isinstance(v, bool): 359 | return v 360 | if v.lower() in ("yes", "true", "t", "y", "1"): 361 | return True 362 | elif v.lower() in ("no", "false", "f", "n", "0"): 363 | return False 364 | else: 365 | raise argparse.ArgumentTypeError("boolean value expected") 366 | -------------------------------------------------------------------------------- /cm/logger.py: -------------------------------------------------------------------------------- 1 | """ 2 | Logger copied from OpenAI baselines to avoid extra RL-based dependencies: 3 | https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/logger.py 4 | """ 5 | 6 | import os 7 | import sys 8 | import shutil 9 | import os.path as osp 10 | import json 11 | import time 12 | import datetime 13 | import tempfile 14 | import warnings 15 | from collections import defaultdict 16 | from contextlib import contextmanager 17 | 18 | DEBUG = 10 19 | INFO = 20 20 | WARN = 30 21 | ERROR = 40 22 | 23 | DISABLED = 50 24 | 25 | 26 | class KVWriter(object): 27 | def writekvs(self, kvs): 28 | raise NotImplementedError 29 | 30 | 31 | class SeqWriter(object): 32 | def writeseq(self, seq): 33 | raise NotImplementedError 34 | 35 | 36 | class HumanOutputFormat(KVWriter, SeqWriter): 37 | def __init__(self, filename_or_file): 38 | if isinstance(filename_or_file, str): 39 | self.file = open(filename_or_file, "wt") 40 | self.own_file = True 41 | else: 42 | assert hasattr(filename_or_file, "read"), ( 43 | "expected file or str, got %s" % filename_or_file 44 | ) 45 | self.file = filename_or_file 46 | self.own_file = False 47 | 48 | def writekvs(self, kvs): 49 | # Create strings for printing 50 | key2str = {} 51 | for (key, val) in sorted(kvs.items()): 52 | if hasattr(val, "__float__"): 53 | valstr = "%-8.3g" % val 54 | else: 55 | valstr = str(val) 56 | key2str[self._truncate(key)] = self._truncate(valstr) 57 | 58 | # Find max widths 59 | if len(key2str) == 0: 60 | print("WARNING: tried to write empty key-value dict") 61 | return 62 | else: 63 | keywidth = max(map(len, key2str.keys())) 64 | valwidth = max(map(len, key2str.values())) 65 | 66 | # Write out the data 67 | dashes = "-" * (keywidth + valwidth + 7) 68 | lines = [dashes] 69 | for (key, val) in sorted(key2str.items(), key=lambda kv: kv[0].lower()): 70 | lines.append( 71 | "| %s%s | %s%s |" 72 | % (key, " " * (keywidth - len(key)), val, " " * (valwidth - len(val))) 73 | ) 74 | lines.append(dashes) 75 | self.file.write("\n".join(lines) + "\n") 76 | 77 | # Flush the output to the file 78 | self.file.flush() 79 | 80 | def _truncate(self, s): 81 | maxlen = 30 82 | return s[: maxlen - 3] + "..." if len(s) > maxlen else s 83 | 84 | def writeseq(self, seq): 85 | seq = list(seq) 86 | for (i, elem) in enumerate(seq): 87 | self.file.write(elem) 88 | if i < len(seq) - 1: # add space unless this is the last one 89 | self.file.write(" ") 90 | self.file.write("\n") 91 | self.file.flush() 92 | 93 | def close(self): 94 | if self.own_file: 95 | self.file.close() 96 | 97 | 98 | class JSONOutputFormat(KVWriter): 99 | def __init__(self, filename): 100 | self.file = open(filename, "wt") 101 | 102 | def writekvs(self, kvs): 103 | for k, v in sorted(kvs.items()): 104 | if hasattr(v, "dtype"): 105 | kvs[k] = float(v) 106 | self.file.write(json.dumps(kvs) + "\n") 107 | self.file.flush() 108 | 109 | def close(self): 110 | self.file.close() 111 | 112 | 113 | class CSVOutputFormat(KVWriter): 114 | def __init__(self, filename): 115 | self.file = open(filename, "w+t") 116 | self.keys = [] 117 | self.sep = "," 118 | 119 | def writekvs(self, kvs): 120 | # Add our current row to the history 121 | extra_keys = list(kvs.keys() - self.keys) 122 | extra_keys.sort() 123 | if extra_keys: 124 | self.keys.extend(extra_keys) 125 | self.file.seek(0) 126 | lines = self.file.readlines() 127 | self.file.seek(0) 128 | for (i, k) in enumerate(self.keys): 129 | if i > 0: 130 | self.file.write(",") 131 | self.file.write(k) 132 | self.file.write("\n") 133 | for line in lines[1:]: 134 | self.file.write(line[:-1]) 135 | self.file.write(self.sep * len(extra_keys)) 136 | self.file.write("\n") 137 | for (i, k) in enumerate(self.keys): 138 | if i > 0: 139 | self.file.write(",") 140 | v = kvs.get(k) 141 | if v is not None: 142 | self.file.write(str(v)) 143 | self.file.write("\n") 144 | self.file.flush() 145 | 146 | def close(self): 147 | self.file.close() 148 | 149 | 150 | class TensorBoardOutputFormat(KVWriter): 151 | """ 152 | Dumps key/value pairs into TensorBoard's numeric format. 153 | """ 154 | 155 | def __init__(self, dir): 156 | os.makedirs(dir, exist_ok=True) 157 | self.dir = dir 158 | self.step = 1 159 | prefix = "events" 160 | path = osp.join(osp.abspath(dir), prefix) 161 | import tensorflow as tf 162 | from tensorflow.python import pywrap_tensorflow 163 | from tensorflow.core.util import event_pb2 164 | from tensorflow.python.util import compat 165 | 166 | self.tf = tf 167 | self.event_pb2 = event_pb2 168 | self.pywrap_tensorflow = pywrap_tensorflow 169 | self.writer = pywrap_tensorflow.EventsWriter(compat.as_bytes(path)) 170 | 171 | def writekvs(self, kvs): 172 | def summary_val(k, v): 173 | kwargs = {"tag": k, "simple_value": float(v)} 174 | return self.tf.Summary.Value(**kwargs) 175 | 176 | summary = self.tf.Summary(value=[summary_val(k, v) for k, v in kvs.items()]) 177 | event = self.event_pb2.Event(wall_time=time.time(), summary=summary) 178 | event.step = ( 179 | self.step 180 | ) # is there any reason why you'd want to specify the step? 181 | self.writer.WriteEvent(event) 182 | self.writer.Flush() 183 | self.step += 1 184 | 185 | def close(self): 186 | if self.writer: 187 | self.writer.Close() 188 | self.writer = None 189 | 190 | 191 | def make_output_format(format, ev_dir, log_suffix=""): 192 | os.makedirs(ev_dir, exist_ok=True) 193 | if format == "stdout": 194 | return HumanOutputFormat(sys.stdout) 195 | elif format == "log": 196 | return HumanOutputFormat(osp.join(ev_dir, "log%s.txt" % log_suffix)) 197 | elif format == "json": 198 | return JSONOutputFormat(osp.join(ev_dir, "progress%s.json" % log_suffix)) 199 | elif format == "csv": 200 | return CSVOutputFormat(osp.join(ev_dir, "progress%s.csv" % log_suffix)) 201 | elif format == "tensorboard": 202 | return TensorBoardOutputFormat(osp.join(ev_dir, "tb%s" % log_suffix)) 203 | else: 204 | raise ValueError("Unknown format specified: %s" % (format,)) 205 | 206 | 207 | # ================================================================ 208 | # API 209 | # ================================================================ 210 | 211 | 212 | def logkv(key, val): 213 | """ 214 | Log a value of some diagnostic 215 | Call this once for each diagnostic quantity, each iteration 216 | If called many times, last value will be used. 217 | """ 218 | get_current().logkv(key, val) 219 | 220 | 221 | def logkv_mean(key, val): 222 | """ 223 | The same as logkv(), but if called many times, values averaged. 224 | """ 225 | get_current().logkv_mean(key, val) 226 | 227 | 228 | def logkvs(d): 229 | """ 230 | Log a dictionary of key-value pairs 231 | """ 232 | for (k, v) in d.items(): 233 | logkv(k, v) 234 | 235 | 236 | def dumpkvs(): 237 | """ 238 | Write all of the diagnostics from the current iteration 239 | """ 240 | return get_current().dumpkvs() 241 | 242 | 243 | def getkvs(): 244 | return get_current().name2val 245 | 246 | 247 | def log(*args, level=INFO): 248 | """ 249 | Write the sequence of args, with no separators, to the console and output files (if you've configured an output file). 250 | """ 251 | get_current().log(*args, level=level) 252 | 253 | 254 | def debug(*args): 255 | log(*args, level=DEBUG) 256 | 257 | 258 | def info(*args): 259 | log(*args, level=INFO) 260 | 261 | 262 | def warn(*args): 263 | log(*args, level=WARN) 264 | 265 | 266 | def error(*args): 267 | log(*args, level=ERROR) 268 | 269 | 270 | def set_level(level): 271 | """ 272 | Set logging threshold on current logger. 273 | """ 274 | get_current().set_level(level) 275 | 276 | 277 | def set_comm(comm): 278 | get_current().set_comm(comm) 279 | 280 | 281 | def get_dir(): 282 | """ 283 | Get directory that log files are being written to. 284 | will be None if there is no output directory (i.e., if you didn't call start) 285 | """ 286 | return get_current().get_dir() 287 | 288 | 289 | record_tabular = logkv 290 | dump_tabular = dumpkvs 291 | 292 | 293 | @contextmanager 294 | def profile_kv(scopename): 295 | logkey = "wait_" + scopename 296 | tstart = time.time() 297 | try: 298 | yield 299 | finally: 300 | get_current().name2val[logkey] += time.time() - tstart 301 | 302 | 303 | def profile(n): 304 | """ 305 | Usage: 306 | @profile("my_func") 307 | def my_func(): code 308 | """ 309 | 310 | def decorator_with_name(func): 311 | def func_wrapper(*args, **kwargs): 312 | with profile_kv(n): 313 | return func(*args, **kwargs) 314 | 315 | return func_wrapper 316 | 317 | return decorator_with_name 318 | 319 | 320 | # ================================================================ 321 | # Backend 322 | # ================================================================ 323 | 324 | 325 | def get_current(): 326 | if Logger.CURRENT is None: 327 | _configure_default_logger() 328 | 329 | return Logger.CURRENT 330 | 331 | 332 | class Logger(object): 333 | DEFAULT = None # A logger with no output files. (See right below class definition) 334 | # So that you can still log to the terminal without setting up any output files 335 | CURRENT = None # Current logger being used by the free functions above 336 | 337 | def __init__(self, dir, output_formats, comm=None): 338 | self.name2val = defaultdict(float) # values this iteration 339 | self.name2cnt = defaultdict(int) 340 | self.level = INFO 341 | self.dir = dir 342 | self.output_formats = output_formats 343 | self.comm = comm 344 | 345 | # Logging API, forwarded 346 | # ---------------------------------------- 347 | def logkv(self, key, val): 348 | self.name2val[key] = val 349 | 350 | def logkv_mean(self, key, val): 351 | oldval, cnt = self.name2val[key], self.name2cnt[key] 352 | self.name2val[key] = oldval * cnt / (cnt + 1) + val / (cnt + 1) 353 | self.name2cnt[key] = cnt + 1 354 | 355 | def dumpkvs(self): 356 | if self.comm is None: 357 | d = self.name2val 358 | else: 359 | d = mpi_weighted_mean( 360 | self.comm, 361 | { 362 | name: (val, self.name2cnt.get(name, 1)) 363 | for (name, val) in self.name2val.items() 364 | }, 365 | ) 366 | if self.comm.rank != 0: 367 | d["dummy"] = 1 # so we don't get a warning about empty dict 368 | out = d.copy() # Return the dict for unit testing purposes 369 | for fmt in self.output_formats: 370 | if isinstance(fmt, KVWriter): 371 | fmt.writekvs(d) 372 | self.name2val.clear() 373 | self.name2cnt.clear() 374 | return out 375 | 376 | def log(self, *args, level=INFO): 377 | if self.level <= level: 378 | self._do_log(args) 379 | 380 | # Configuration 381 | # ---------------------------------------- 382 | def set_level(self, level): 383 | self.level = level 384 | 385 | def set_comm(self, comm): 386 | self.comm = comm 387 | 388 | def get_dir(self): 389 | return self.dir 390 | 391 | def close(self): 392 | for fmt in self.output_formats: 393 | fmt.close() 394 | 395 | # Misc 396 | # ---------------------------------------- 397 | def _do_log(self, args): 398 | for fmt in self.output_formats: 399 | if isinstance(fmt, SeqWriter): 400 | fmt.writeseq(map(str, args)) 401 | 402 | 403 | def get_rank_without_mpi_import(): 404 | # check environment variables here instead of importing mpi4py 405 | # to avoid calling MPI_Init() when this module is imported 406 | for varname in ["PMI_RANK", "OMPI_COMM_WORLD_RANK"]: 407 | if varname in os.environ: 408 | return int(os.environ[varname]) 409 | return 0 410 | 411 | 412 | def mpi_weighted_mean(comm, local_name2valcount): 413 | """ 414 | Copied from: https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/common/mpi_util.py#L110 415 | Perform a weighted average over dicts that are each on a different node 416 | Input: local_name2valcount: dict mapping key -> (value, count) 417 | Returns: key -> mean 418 | """ 419 | all_name2valcount = comm.gather(local_name2valcount) 420 | if comm.rank == 0: 421 | name2sum = defaultdict(float) 422 | name2count = defaultdict(float) 423 | for n2vc in all_name2valcount: 424 | for (name, (val, count)) in n2vc.items(): 425 | try: 426 | val = float(val) 427 | except ValueError: 428 | if comm.rank == 0: 429 | warnings.warn( 430 | "WARNING: tried to compute mean on non-float {}={}".format( 431 | name, val 432 | ) 433 | ) 434 | else: 435 | name2sum[name] += val * count 436 | name2count[name] += count 437 | return {name: name2sum[name] / name2count[name] for name in name2sum} 438 | else: 439 | return {} 440 | 441 | 442 | def configure(dir=None, format_strs=None, comm=None, log_suffix=""): 443 | """ 444 | If comm is provided, average all numerical stats across that comm 445 | """ 446 | if dir is None: 447 | dir = os.getenv("OPENAI_LOGDIR") 448 | if dir is None: 449 | # dir = osp.join( 450 | # tempfile.gettempdir(), 451 | # datetime.datetime.now().strftime("openai-%Y-%m-%d-%H-%M-%S-%f"), 452 | # ) 453 | dir = osp.join( 454 | "../experiment", 455 | datetime.datetime.now().strftime("openai-%Y-%m-%d-%H-%M-%S-%f"), 456 | ) 457 | assert isinstance(dir, str) 458 | dir = os.path.expanduser(dir) 459 | os.makedirs(os.path.expanduser(dir), exist_ok=True) 460 | 461 | rank = get_rank_without_mpi_import() 462 | if rank > 0: 463 | log_suffix = log_suffix + "-rank%03i" % rank 464 | 465 | if format_strs is None: 466 | if rank == 0: 467 | format_strs = os.getenv("OPENAI_LOG_FORMAT", "stdout,log,csv").split(",") 468 | else: 469 | format_strs = os.getenv("OPENAI_LOG_FORMAT_MPI", "log").split(",") 470 | format_strs = filter(None, format_strs) 471 | output_formats = [make_output_format(f, dir, log_suffix) for f in format_strs] 472 | 473 | Logger.CURRENT = Logger(dir=dir, output_formats=output_formats, comm=comm) 474 | if output_formats: 475 | log("Logging to %s" % dir) 476 | 477 | 478 | def _configure_default_logger(): 479 | configure() 480 | Logger.DEFAULT = Logger.CURRENT 481 | 482 | 483 | def reset(): 484 | if Logger.CURRENT is not Logger.DEFAULT: 485 | Logger.CURRENT.close() 486 | Logger.CURRENT = Logger.DEFAULT 487 | log("Reset logger") 488 | 489 | 490 | @contextmanager 491 | def scoped_configure(dir=None, format_strs=None, comm=None): 492 | prevlogger = Logger.CURRENT 493 | configure(dir=dir, format_strs=format_strs, comm=comm) 494 | try: 495 | yield 496 | finally: 497 | Logger.CURRENT.close() 498 | Logger.CURRENT = prevlogger 499 | 500 | -------------------------------------------------------------------------------- /cm/train_util.py: -------------------------------------------------------------------------------- 1 | import copy 2 | import functools 3 | import os 4 | 5 | import blobfile as bf 6 | import torch as th 7 | import torch.distributed as dist 8 | from torch.nn.parallel.distributed import DistributedDataParallel as DDP 9 | from torch.optim import RAdam, Adam 10 | 11 | from . import dist_util, logger 12 | from .fp16_util import MixedPrecisionTrainer 13 | from .nn import update_ema 14 | from .resample import LossAwareSampler, UniformSampler 15 | 16 | from .fp16_util import ( 17 | get_param_groups_and_shapes, 18 | make_master_params, 19 | master_params_to_model_params, 20 | ) 21 | import numpy as np 22 | import torchvision 23 | import time 24 | import datetime 25 | import copy 26 | 27 | # For ImageNet experiments, this was a good default value. 28 | # We found that the lg_loss_scale quickly climbed to 29 | # 20-21 within the first ~1K steps of training. 30 | INITIAL_LOG_LOSS_SCALE = 20.0 31 | 32 | def sec2str(seconds): 33 | days = seconds // (24 * 3600) 34 | seconds = seconds % (24 * 3600) 35 | hour = seconds // 3600 36 | seconds %= 3600 37 | minutes = seconds // 60 38 | seconds %= 60 39 | 40 | return "%d:%d:%02d:%02d" % (days, hour, minutes, seconds) 41 | 42 | class TrainLoop: 43 | def __init__( 44 | self, 45 | *, 46 | model, 47 | diffusion, 48 | data, 49 | batch_size, 50 | microbatch, 51 | lr, 52 | ema_rate, 53 | log_interval, 54 | save_interval, 55 | resume_checkpoint, 56 | use_fp16=False, 57 | fp16_scale_growth=1e-3, 58 | schedule_sampler=None, 59 | weight_decay=0.0, 60 | lr_anneal_steps=0, 61 | optimizer='radam', 62 | opt_eps=1e-8, 63 | ): 64 | self.model = model 65 | self.diffusion = diffusion 66 | self.data = data 67 | self.batch_size = batch_size 68 | self.microbatch = microbatch if microbatch > 0 else batch_size 69 | self.lr = lr 70 | self.ema_rate = ( 71 | [ema_rate] 72 | if isinstance(ema_rate, float) 73 | else [float(x) for x in ema_rate.split(",")] 74 | ) 75 | self.log_interval = log_interval 76 | self.save_interval = save_interval 77 | self.resume_checkpoint = resume_checkpoint 78 | self.use_fp16 = use_fp16 79 | self.fp16_scale_growth = fp16_scale_growth 80 | self.schedule_sampler = schedule_sampler or UniformSampler(diffusion) 81 | self.weight_decay = weight_decay 82 | self.lr_anneal_steps = lr_anneal_steps 83 | 84 | self.step = 0 85 | self.resume_step = 0 86 | self.global_batch = self.batch_size * dist.get_world_size() 87 | 88 | self.sync_cuda = th.cuda.is_available() 89 | 90 | self._load_and_sync_parameters() 91 | self.mp_trainer = MixedPrecisionTrainer( 92 | model=self.model, 93 | use_fp16=self.use_fp16, 94 | fp16_scale_growth=fp16_scale_growth, 95 | ) 96 | 97 | training_params = self.mp_trainer.master_params 98 | 99 | if optimizer == 'radam': 100 | logger.log(f"use optimizer RAdam, lr: {self.lr}, weight_decay: {self.weight_decay}, eps: {opt_eps}") 101 | self.opt = RAdam( 102 | training_params , lr=self.lr, weight_decay=self.weight_decay, eps=opt_eps 103 | ) 104 | elif optimizer == 'adam': 105 | logger.log(f"use optimizer Adam, lr: {self.lr}, weight_decay: {self.weight_decay}, eps: {opt_eps}") 106 | self.opt = Adam( 107 | training_params, lr=self.lr, weight_decay=self.weight_decay, eps=opt_eps 108 | ) 109 | else: 110 | raise ValueError(f"Unsporrt optimizer {optimizer}") 111 | 112 | if self.resume_step: 113 | self._load_optimizer_state() 114 | # Model was resumed, either due to a restart or a checkpoint 115 | # being specified at the command line. 116 | self.ema_params = [ 117 | self._load_ema_parameters(rate) for rate in self.ema_rate 118 | ] 119 | else: 120 | self.ema_params = [ 121 | copy.deepcopy(self.mp_trainer.master_params) 122 | for _ in range(len(self.ema_rate)) 123 | ] 124 | 125 | if th.cuda.is_available(): 126 | self.use_ddp = True 127 | self.ddp_model = DDP( 128 | self.model, 129 | device_ids=[dist_util.dev()], 130 | output_device=dist_util.dev(), 131 | broadcast_buffers=False, 132 | bucket_cap_mb=128, 133 | find_unused_parameters=False, 134 | ) 135 | 136 | else: 137 | if dist.get_world_size() > 1: 138 | logger.warn( 139 | "Distributed training requires CUDA. " 140 | "Gradients will not be synchronized properly!" 141 | ) 142 | self.use_ddp = False 143 | self.ddp_model = self.model 144 | 145 | self.step = self.resume_step 146 | 147 | if dist.get_rank() == 0: 148 | code_save_dir = bf.join(get_blob_logdir(), "code") 149 | # try: 150 | import shutil 151 | shutil.copytree("./scripts", bf.join(code_save_dir, 'scripts') , dirs_exist_ok=True) 152 | shutil.copytree("./cm", bf.join(code_save_dir, 'cm') , dirs_exist_ok=True) 153 | logger.log(f"Save the code for current run in {code_save_dir}") 154 | # except Exception as err: 155 | # logger.log(f"Fail to save the code. {err}") 156 | # logger.log(f"Continue without saving the code...") 157 | 158 | 159 | 160 | def _load_and_sync_parameters(self): 161 | resume_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 162 | 163 | if resume_checkpoint: 164 | self.resume_step = parse_resume_step_from_filename(resume_checkpoint) 165 | if dist.get_rank() == 0: 166 | logger.log(f"loading model from checkpoint: {resume_checkpoint}...") 167 | self.model.load_state_dict( 168 | dist_util.load_state_dict( 169 | resume_checkpoint, map_location=dist_util.dev() 170 | ), 171 | ) 172 | 173 | dist_util.sync_params(self.model.parameters()) 174 | dist_util.sync_params(self.model.buffers()) 175 | 176 | def _load_ema_parameters(self, rate): 177 | ema_params = copy.deepcopy(self.mp_trainer.master_params) 178 | 179 | main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 180 | ema_checkpoint = find_ema_checkpoint(main_checkpoint, self.resume_step, rate) 181 | if ema_checkpoint: 182 | if dist.get_rank() == 0: 183 | logger.log(f"loading EMA from checkpoint: {ema_checkpoint}...") 184 | state_dict = dist_util.load_state_dict( 185 | ema_checkpoint, map_location=dist_util.dev() 186 | ) 187 | ema_params = self.mp_trainer.state_dict_to_master_params(state_dict) 188 | 189 | dist_util.sync_params(ema_params) 190 | return ema_params 191 | 192 | def _load_optimizer_state(self): 193 | main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 194 | opt_checkpoint = bf.join( 195 | bf.dirname(main_checkpoint), f"opt{self.resume_step:06}.pt" 196 | ) 197 | if bf.exists(opt_checkpoint): 198 | logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}") 199 | state_dict = dist_util.load_state_dict( 200 | opt_checkpoint, map_location=dist_util.dev() 201 | ) 202 | self.opt.load_state_dict(state_dict) 203 | 204 | def run_loop(self): 205 | while not self.lr_anneal_steps or self.step < self.lr_anneal_steps: 206 | batch, cond = next(self.data) 207 | self.run_step(batch, cond) 208 | if self.step % self.log_interval == 0: 209 | logger.dumpkvs() 210 | if self.step % self.save_interval == 0: 211 | self.save() 212 | # Run for a finite amount of time in integration tests. 213 | if os.environ.get("DIFFUSION_TRAINING_TEST", "") and self.step > 0: 214 | return 215 | # Save the last checkpoint if it wasn't already saved. 216 | if (self.step - 1) % self.save_interval != 0: 217 | self.save() 218 | 219 | def run_step(self, batch, cond): 220 | self.forward_backward(batch, cond) 221 | took_step = self.mp_trainer.optimize(self.opt) 222 | if took_step: 223 | self.step += 1 224 | self._update_ema() 225 | self._anneal_lr() 226 | self.log_step() 227 | 228 | def forward_backward(self, batch, cond): 229 | self.mp_trainer.zero_grad() 230 | for i in range(0, batch.shape[0], self.microbatch): 231 | micro = batch[i : i + self.microbatch].to(dist_util.dev()) 232 | micro_cond = { 233 | k: v[i : i + self.microbatch].to(dist_util.dev()) 234 | for k, v in cond.items() 235 | } 236 | last_batch = (i + self.microbatch) >= batch.shape[0] 237 | t, weights = self.schedule_sampler.sample(micro.shape[0], dist_util.dev()) 238 | 239 | compute_losses = functools.partial( 240 | self.diffusion.training_losses, 241 | self.ddp_model, 242 | micro, 243 | t, 244 | model_kwargs=micro_cond, 245 | ) 246 | 247 | if last_batch or not self.use_ddp: 248 | losses = compute_losses() 249 | else: 250 | with self.ddp_model.no_sync(): 251 | losses = compute_losses() 252 | 253 | if isinstance(self.schedule_sampler, LossAwareSampler): 254 | self.schedule_sampler.update_with_local_losses( 255 | t, losses["loss"].detach() 256 | ) 257 | 258 | loss = (losses["loss"] * weights).mean() 259 | log_loss_dict( 260 | self.diffusion, t, {k: v * weights for k, v in losses.items()} 261 | ) 262 | self.mp_trainer.backward(loss) 263 | 264 | def _update_ema(self): 265 | for rate, params in zip(self.ema_rate, self.ema_params): 266 | update_ema(params, self.mp_trainer.master_params, rate=rate) 267 | 268 | 269 | def _anneal_lr(self): 270 | if not self.lr_anneal_steps: 271 | return 272 | frac_done = (self.step + self.resume_step) / self.lr_anneal_steps 273 | lr = self.lr * (1 - frac_done) 274 | for param_group in self.opt.param_groups: 275 | param_group["lr"] = lr 276 | 277 | def log_step(self): 278 | logger.logkv("step", self.step + self.resume_step) 279 | logger.logkv("samples", (self.step + self.resume_step + 1) * self.global_batch) 280 | 281 | def save(self): 282 | def save_checkpoint(rate, params): 283 | state_dict = self.mp_trainer.master_params_to_state_dict(params) 284 | if dist.get_rank() == 0: 285 | logger.log(f"saving model {rate}...") 286 | if not rate: 287 | filename = f"model{(self.step+self.resume_step):06d}.pt" 288 | else: 289 | filename = f"ema_{rate}_{(self.step+self.resume_step):06d}.pt" 290 | with bf.BlobFile(bf.join(get_blob_logdir(), filename), "wb") as f: 291 | th.save(state_dict, f) 292 | 293 | for rate, params in zip(self.ema_rate, self.ema_params): 294 | save_checkpoint(rate, params) 295 | 296 | if dist.get_rank() == 0: 297 | with bf.BlobFile( 298 | bf.join(get_blob_logdir(), f"opt{(self.step+self.resume_step):06d}.pt"), 299 | "wb", 300 | ) as f: 301 | th.save(self.opt.state_dict(), f) 302 | 303 | # Save model parameters last to prevent race conditions where a restart 304 | # loads model at step N, but opt/ema state isn't saved for step N. 305 | save_checkpoint(0, self.mp_trainer.master_params) 306 | dist.barrier() 307 | 308 | 309 | 310 | class ODETrainLoop(TrainLoop): 311 | def __init__( 312 | self, 313 | *, 314 | teacher_model, 315 | teacher_diffusion, 316 | training_mode, 317 | ema_scale_fn, 318 | total_training_steps, 319 | methodology, 320 | eval_interval, 321 | **kwargs, 322 | ): 323 | super().__init__(**kwargs) 324 | self.training_mode = training_mode 325 | self.ema_scale_fn = ema_scale_fn 326 | self.teacher_model = teacher_model 327 | self.teacher_diffusion = teacher_diffusion 328 | self.total_training_steps = total_training_steps 329 | self.methodology = methodology 330 | self.eval_interval = eval_interval 331 | 332 | 333 | 334 | if teacher_model: 335 | self._load_and_sync_teacher_parameters() 336 | self.teacher_model.requires_grad_(False) 337 | self.teacher_model.eval() 338 | for param in self.teacher_model.parameters(): 339 | param.requires_grad_(False) 340 | 341 | self.global_step = self.step 342 | if training_mode == "progdist": 343 | _, scale = ema_scale_fn(self.global_step) 344 | if scale == 1 or scale == 2: 345 | _, start_scale = ema_scale_fn(0) 346 | n_normal_steps = int(np.log2(start_scale // 2)) * self.lr_anneal_steps 347 | step = self.global_step - n_normal_steps 348 | if step != 0: 349 | self.lr_anneal_steps *= 2 350 | self.step = step % self.lr_anneal_steps 351 | else: 352 | self.step = 0 353 | else: 354 | self.step = self.global_step % self.lr_anneal_steps 355 | 356 | 357 | 358 | def _load_and_sync_teacher_parameters(self): 359 | resume_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 360 | if resume_checkpoint: 361 | path, name = os.path.split(resume_checkpoint) 362 | teacher_name = name.replace("model", "teacher_model") 363 | resume_teacher_checkpoint = os.path.join(path, teacher_name) 364 | 365 | if bf.exists(resume_teacher_checkpoint) and dist.get_rank() == 0: 366 | logger.log( 367 | "loading model from checkpoint: {resume_teacher_checkpoint}..." 368 | ) 369 | self.teacher_model.load_state_dict( 370 | dist_util.load_state_dict( 371 | resume_teacher_checkpoint, map_location=dist_util.dev() 372 | ), 373 | ) 374 | 375 | dist_util.sync_params(self.teacher_model.parameters()) 376 | dist_util.sync_params(self.teacher_model.buffers()) 377 | 378 | def run_loop(self): 379 | saved = False 380 | from cm.unet import UNetModel 381 | self.model: UNetModel 382 | self.image_size = (self.model.in_channels, self.model.image_size, self.model.image_size) 383 | tic_time = time.time() 384 | begining_time = time.time() 385 | 386 | while ( 387 | not self.lr_anneal_steps 388 | or self.step < self.lr_anneal_steps 389 | or self.global_step < self.total_training_steps 390 | ): 391 | 392 | device = dist_util.dev() 393 | 394 | ########################################################################## 395 | # random noise 396 | ########################################################################## 397 | init_condition = th.randn((self.batch_size, *self.image_size), device=device) 398 | training_class_labels = None 399 | if self.model.num_classes is not None: 400 | training_class_labels = th.randint(0, self.model.num_classes, (self.batch_size,), device=device) 401 | 402 | ########################################################################## 403 | 404 | # construct the dataset 405 | 406 | batch = { 407 | "noise": init_condition * 80., 408 | } 409 | cond = { 410 | "noise_cond":training_class_labels, 411 | } 412 | 413 | self.run_step(batch, cond) 414 | saved = False 415 | if ( 416 | self.global_step 417 | and self.save_interval != -1 418 | and self.global_step % self.save_interval == 0 419 | ): 420 | self.save() 421 | saved = True 422 | th.cuda.empty_cache() 423 | # Run for a finite amount of time in integration tests. 424 | if os.environ.get("DIFFUSION_TRAINING_TEST", "") and self.step > 0: 425 | return 426 | 427 | if self.global_step % self.log_interval == 0: 428 | logger.dumpkvs() 429 | 430 | if dist.get_rank() == 0: 431 | current_time = datetime.datetime.now() 432 | t_dur = time.time() - tic_time 433 | t_spend = time.time() - begining_time 434 | t_eta = t_dur*((self.total_training_steps-self.global_step)//(self.log_interval)) 435 | t_total = t_eta + t_spend 436 | 437 | logger.log(f"{t_dur} s/10iter at {current_time}. ETA: {sec2str(t_eta)}, Spend: {sec2str(t_spend)}, Total: {sec2str(t_total)}") 438 | 439 | tic_time = time.time() 440 | 441 | # Save the last checkpoint if it wasn't already saved. 442 | if not saved: 443 | self.save() 444 | 445 | 446 | def run_step(self, batch, cond): 447 | self.forward_backward(batch, cond) 448 | took_step = self.mp_trainer.optimize(self.opt) 449 | if took_step: 450 | self._update_ema() 451 | 452 | self.step += 1 453 | self.global_step += 1 454 | 455 | self._anneal_lr() 456 | self.log_step() 457 | 458 | 459 | 460 | 461 | def forward_backward(self, batch, cond): 462 | self.mp_trainer.zero_grad() 463 | batch_size = batch['noise'].shape[0] 464 | for i in range(0, batch_size, self.microbatch): 465 | micro = {} 466 | for k, v in batch.items(): 467 | micro[k] = v[i : i + self.microbatch].to(dist_util.dev()) 468 | 469 | micro_cond = None 470 | if self.model.num_classes: 471 | if self.training_mode in ["one_shot_pinn_edm_edm_randominit"]: 472 | micro_cond = { 473 | "noise_cond": cond["noise_cond"][i : i + self.microbatch].to(dist_util.dev()), 474 | "img_cond": cond["img_cond"][i : i + self.microbatch].to(dist_util.dev()), 475 | } 476 | else: 477 | micro_cond = {"noise_cond": cond["noise_cond"][i : i + self.microbatch].to(dist_util.dev())} 478 | 479 | 480 | last_batch = (i + self.microbatch) >= batch_size 481 | t, weights = self.schedule_sampler.sample(micro['noise'].shape[0], dist_util.dev()) 482 | 483 | ema, num_scales = self.ema_scale_fn(self.global_step) 484 | 485 | compute_losses = functools.partial( 486 | self.diffusion.ode_losses, 487 | self.ddp_model, 488 | micro, 489 | num_scales, 490 | teacher_model=self.teacher_model, 491 | teacher_diffusion=self.teacher_diffusion, 492 | model_kwargs=micro_cond, 493 | current_step=self.global_step, 494 | ) 495 | 496 | if last_batch or not self.use_ddp: 497 | losses = compute_losses() 498 | else: 499 | with self.ddp_model.no_sync(): 500 | losses = compute_losses() 501 | 502 | if isinstance(self.schedule_sampler, LossAwareSampler): 503 | self.schedule_sampler.update_with_local_losses( 504 | t, losses["loss"].detach() 505 | ) 506 | 507 | loss = (losses["loss"]).mean() 508 | 509 | log_loss_dict( 510 | self.diffusion, t, {k: v for k, v in losses.items()} 511 | ) 512 | self.mp_trainer.backward(loss) 513 | 514 | def save(self): 515 | import blobfile as bf 516 | 517 | step = self.global_step 518 | 519 | def save_checkpoint(rate, params): 520 | state_dict = self.mp_trainer.master_params_to_state_dict(params) 521 | if dist.get_rank() == 0: 522 | logger.log(f"saving model {rate}...") 523 | if not rate: 524 | filename = f"model{step:06d}.pt" 525 | else: 526 | filename = f"ema_{rate}_{step:06d}.pt" 527 | with bf.BlobFile(bf.join(get_blob_logdir(), filename), "wb") as f: 528 | th.save(state_dict, f) 529 | 530 | for rate, params in zip(self.ema_rate, self.ema_params): 531 | save_checkpoint(rate, params) 532 | 533 | logger.log("saving optimizer state...") 534 | if dist.get_rank() == 0: 535 | with bf.BlobFile( 536 | bf.join(get_blob_logdir(), f"opt{step:06d}.pt"), 537 | "wb", 538 | ) as f: 539 | th.save(self.opt.state_dict(), f) 540 | 541 | if dist.get_rank() == 0: 542 | 543 | if self.teacher_model and self.training_mode == "progdist": 544 | logger.log("saving teacher model state") 545 | filename = f"teacher_model{step:06d}.pt" 546 | with bf.BlobFile(bf.join(get_blob_logdir(), filename), "wb") as f: 547 | th.save(self.teacher_model.state_dict(), f) 548 | 549 | # Save model parameters last to prevent race conditions where a restart 550 | # loads model at step N, but opt/ema state isn't saved for step N. 551 | save_checkpoint(0, self.mp_trainer.master_params) 552 | 553 | dist.barrier() 554 | 555 | def log_step(self): 556 | step = self.global_step 557 | logger.logkv("step", step) 558 | logger.logkv("samples", (step + 1) * self.global_batch) 559 | 560 | 561 | 562 | def parse_resume_step_from_filename(filename): 563 | """ 564 | Parse filenames of the form path/to/modelNNNNNN.pt, where NNNNNN is the 565 | checkpoint's number of steps. 566 | """ 567 | split = filename.split("model") 568 | if len(split) < 2: 569 | return 0 570 | split1 = split[-1].split(".")[0] 571 | try: 572 | return int(split1) 573 | except ValueError: 574 | return 0 575 | 576 | 577 | def get_blob_logdir(): 578 | # You can change this to be a separate path to save checkpoints to 579 | # a blobstore or some external drive. 580 | return logger.get_dir() 581 | 582 | 583 | def find_resume_checkpoint(): 584 | # On your infrastructure, you may want to override this to automatically 585 | # discover the latest checkpoint on your blob storage, etc. 586 | return None 587 | 588 | 589 | def find_ema_checkpoint(main_checkpoint, step, rate): 590 | if main_checkpoint is None: 591 | return None 592 | filename = f"ema_{rate}_{(step):06d}.pt" 593 | path = bf.join(bf.dirname(main_checkpoint), filename) 594 | if bf.exists(path): 595 | return path 596 | return None 597 | 598 | 599 | def log_loss_dict(diffusion, ts, losses): 600 | for key, values in losses.items(): 601 | logger.logkv_mean(key, values.mean().item()) 602 | # Log the quantiles (four quartiles, in particular). 603 | for sub_t, sub_loss in zip(ts.cpu().numpy(), values.detach().cpu().numpy()): 604 | quartile = int(4 * sub_t / diffusion.num_timesteps) 605 | logger.logkv_mean(f"{key}_q{quartile}", sub_loss) 606 | -------------------------------------------------------------------------------- /cm/unet.py: -------------------------------------------------------------------------------- 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 | checkpoint, 13 | conv_nd, 14 | linear, 15 | avg_pool_nd, 16 | zero_module, 17 | normalization, 18 | timestep_embedding, 19 | ) 20 | 21 | 22 | class AttentionPool2d(nn.Module): 23 | """ 24 | Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py 25 | """ 26 | 27 | def __init__( 28 | self, 29 | spacial_dim: int, 30 | embed_dim: int, 31 | num_heads_channels: int, 32 | output_dim: int = None, 33 | ): 34 | super().__init__() 35 | self.positional_embedding = nn.Parameter( 36 | th.randn(embed_dim, spacial_dim**2 + 1) / embed_dim**0.5 37 | ) 38 | self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1) 39 | self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1) 40 | self.num_heads = embed_dim // num_heads_channels 41 | self.attention = QKVAttention(self.num_heads) 42 | 43 | def forward(self, x): 44 | b, c, *_spatial = x.shape 45 | x = x.reshape(b, c, -1) # NC(HW) 46 | x = th.cat([x.mean(dim=-1, keepdim=True), x], dim=-1) # NC(HW+1) 47 | x = x + self.positional_embedding[None, :, :].to(x.dtype) # NC(HW+1) 48 | x = self.qkv_proj(x) 49 | x = self.attention(x) 50 | x = self.c_proj(x) 51 | return x[:, :, 0] 52 | 53 | 54 | class TimestepBlock(nn.Module): 55 | """ 56 | Any module where forward() takes timestep embeddings as a second argument. 57 | """ 58 | 59 | @abstractmethod 60 | def forward(self, x, emb): 61 | """ 62 | Apply the module to `x` given `emb` timestep embeddings. 63 | """ 64 | 65 | 66 | class TimestepEmbedSequential(nn.Sequential, TimestepBlock): 67 | """ 68 | A sequential module that passes timestep embeddings to the children that 69 | support it as an extra input. 70 | """ 71 | 72 | def forward(self, x, emb): 73 | for layer in self: 74 | if isinstance(layer, TimestepBlock): 75 | x = layer(x, emb) 76 | else: 77 | x = layer(x) 78 | return x 79 | 80 | 81 | class Upsample(nn.Module): 82 | """ 83 | An upsampling 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 | upsampling occurs in the inner-two dimensions. 89 | """ 90 | 91 | def __init__(self, channels, use_conv, dims=2, out_channels=None): 92 | super().__init__() 93 | self.channels = channels 94 | self.out_channels = out_channels or channels 95 | self.use_conv = use_conv 96 | self.dims = dims 97 | if use_conv: 98 | self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=1) 99 | 100 | def forward(self, x): 101 | assert x.shape[1] == self.channels 102 | if self.dims == 3: 103 | x = F.interpolate( 104 | x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest" 105 | ) 106 | else: 107 | x = F.interpolate(x, scale_factor=2, mode="nearest") 108 | if self.use_conv: 109 | x = self.conv(x) 110 | return x 111 | 112 | 113 | class Downsample(nn.Module): 114 | """ 115 | A downsampling layer with an optional convolution. 116 | 117 | :param channels: channels in the inputs and outputs. 118 | :param use_conv: a bool determining if a convolution is applied. 119 | :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then 120 | downsampling occurs in the inner-two dimensions. 121 | """ 122 | 123 | def __init__(self, channels, use_conv, dims=2, out_channels=None): 124 | super().__init__() 125 | self.channels = channels 126 | self.out_channels = out_channels or channels 127 | self.use_conv = use_conv 128 | self.dims = dims 129 | stride = 2 if dims != 3 else (1, 2, 2) 130 | if use_conv: 131 | self.op = conv_nd( 132 | dims, self.channels, self.out_channels, 3, stride=stride, padding=1 133 | ) 134 | else: 135 | assert self.channels == self.out_channels 136 | self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride) 137 | 138 | def forward(self, x): 139 | assert x.shape[1] == self.channels 140 | return self.op(x) 141 | 142 | 143 | class ResBlock(TimestepBlock): 144 | """ 145 | A residual block that can optionally change the number of channels. 146 | 147 | :param channels: the number of input channels. 148 | :param emb_channels: the number of timestep embedding channels. 149 | :param dropout: the rate of dropout. 150 | :param out_channels: if specified, the number of out channels. 151 | :param use_conv: if True and out_channels is specified, use a spatial 152 | convolution instead of a smaller 1x1 convolution to change the 153 | channels in the skip connection. 154 | :param dims: determines if the signal is 1D, 2D, or 3D. 155 | :param use_checkpoint: if True, use gradient checkpointing on this module. 156 | :param up: if True, use this block for upsampling. 157 | :param down: if True, use this block for downsampling. 158 | """ 159 | 160 | def __init__( 161 | self, 162 | channels, 163 | emb_channels, 164 | dropout, 165 | out_channels=None, 166 | use_conv=False, 167 | use_scale_shift_norm=False, 168 | dims=2, 169 | use_checkpoint=False, 170 | up=False, 171 | down=False, 172 | ): 173 | super().__init__() 174 | self.channels = channels 175 | self.emb_channels = emb_channels 176 | self.dropout = dropout 177 | self.out_channels = out_channels or channels 178 | self.use_conv = use_conv 179 | self.use_checkpoint = use_checkpoint 180 | self.use_scale_shift_norm = use_scale_shift_norm 181 | 182 | self.in_layers = nn.Sequential( 183 | normalization(channels), 184 | nn.SiLU(), 185 | conv_nd(dims, channels, self.out_channels, 3, padding=1), 186 | ) 187 | 188 | self.updown = up or down 189 | 190 | if up: 191 | self.h_upd = Upsample(channels, False, dims) 192 | self.x_upd = Upsample(channels, False, dims) 193 | elif down: 194 | self.h_upd = Downsample(channels, False, dims) 195 | self.x_upd = Downsample(channels, False, dims) 196 | else: 197 | self.h_upd = self.x_upd = nn.Identity() 198 | 199 | self.emb_layers = nn.Sequential( 200 | nn.SiLU(), 201 | linear( 202 | emb_channels, 203 | 2 * self.out_channels if use_scale_shift_norm else self.out_channels, 204 | ), 205 | ) 206 | self.out_layers = nn.Sequential( 207 | normalization(self.out_channels), 208 | nn.SiLU(), 209 | nn.Dropout(p=dropout), 210 | zero_module( 211 | conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1) 212 | ), 213 | ) 214 | 215 | if self.out_channels == channels: 216 | self.skip_connection = nn.Identity() 217 | elif use_conv: 218 | self.skip_connection = conv_nd( 219 | dims, channels, self.out_channels, 3, padding=1 220 | ) 221 | else: 222 | self.skip_connection = conv_nd(dims, channels, self.out_channels, 1) 223 | 224 | def forward(self, x, emb): 225 | """ 226 | Apply the block to a Tensor, conditioned on a timestep embedding. 227 | 228 | :param x: an [N x C x ...] Tensor of features. 229 | :param emb: an [N x emb_channels] Tensor of timestep embeddings. 230 | :return: an [N x C x ...] Tensor of outputs. 231 | """ 232 | return checkpoint( 233 | self._forward, (x, emb), self.parameters(), self.use_checkpoint 234 | ) 235 | 236 | def _forward(self, x, emb): 237 | if self.updown: 238 | in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1] 239 | h = in_rest(x) 240 | h = self.h_upd(h) 241 | x = self.x_upd(x) 242 | h = in_conv(h) 243 | else: 244 | h = self.in_layers(x) 245 | emb_out = self.emb_layers(emb).type(h.dtype) 246 | while len(emb_out.shape) < len(h.shape): 247 | emb_out = emb_out[..., None] 248 | if self.use_scale_shift_norm: 249 | out_norm, out_rest = self.out_layers[0], self.out_layers[1:] 250 | scale, shift = th.chunk(emb_out, 2, dim=1) 251 | h = out_norm(h) * (1 + scale) + shift 252 | h = out_rest(h) 253 | else: 254 | h = h + emb_out 255 | h = self.out_layers(h) 256 | return self.skip_connection(x) + h 257 | 258 | 259 | class AttentionBlock(nn.Module): 260 | """ 261 | An attention block that allows spatial positions to attend to each other. 262 | 263 | Originally ported from here, but adapted to the N-d case. 264 | https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66. 265 | """ 266 | 267 | def __init__( 268 | self, 269 | channels, 270 | num_heads=1, 271 | num_head_channels=-1, 272 | use_checkpoint=False, 273 | attention_type="flash", 274 | # attention_type="normal", 275 | encoder_channels=None, 276 | dims=2, 277 | channels_last=False, 278 | use_new_attention_order=False, 279 | ): 280 | super().__init__() 281 | self.channels = channels 282 | if num_head_channels == -1: 283 | self.num_heads = num_heads 284 | else: 285 | assert ( 286 | channels % num_head_channels == 0 287 | ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}" 288 | self.num_heads = channels // num_head_channels 289 | self.use_checkpoint = use_checkpoint 290 | self.norm = normalization(channels) 291 | self.qkv = conv_nd(dims, channels, channels * 3, 1) 292 | self.attention_type = attention_type 293 | if attention_type == "flash": 294 | self.attention = QKVFlashAttention(channels, self.num_heads) 295 | else: 296 | # split heads before split qkv 297 | self.attention = QKVAttentionLegacy(self.num_heads) 298 | 299 | self.use_attention_checkpoint = not ( 300 | self.use_checkpoint or self.attention_type == "flash" 301 | ) 302 | if encoder_channels is not None: 303 | assert attention_type != "flash" 304 | self.encoder_kv = conv_nd(1, encoder_channels, channels * 2, 1) 305 | self.proj_out = zero_module(conv_nd(dims, channels, channels, 1)) 306 | 307 | def forward(self, x, encoder_out=None): 308 | if encoder_out is None: 309 | return checkpoint( 310 | self._forward, (x,), self.parameters(), self.use_checkpoint 311 | ) 312 | else: 313 | return checkpoint( 314 | self._forward, (x, encoder_out), self.parameters(), self.use_checkpoint 315 | ) 316 | 317 | def _forward(self, x, encoder_out=None): 318 | b, _, *spatial = x.shape 319 | qkv = self.qkv(self.norm(x)).view(b, -1, np.prod(spatial)) 320 | if encoder_out is not None: 321 | encoder_out = self.encoder_kv(encoder_out) 322 | h = checkpoint( 323 | self.attention, (qkv, encoder_out), (), self.use_attention_checkpoint 324 | ) 325 | else: 326 | h = checkpoint(self.attention, (qkv,), (), self.use_attention_checkpoint) 327 | h = h.view(b, -1, *spatial) 328 | h = self.proj_out(h) 329 | return x + h 330 | 331 | 332 | class QKVFlashAttention(nn.Module): 333 | def __init__( 334 | self, 335 | embed_dim, 336 | num_heads, 337 | batch_first=True, 338 | attention_dropout=0.0, 339 | causal=False, 340 | device=None, 341 | dtype=None, 342 | **kwargs, 343 | ) -> None: 344 | from einops import rearrange 345 | # from flash_attn.flash_attention import FlashAttention 346 | from flash_attn import flash_attn_qkvpacked_func 347 | 348 | assert batch_first 349 | factory_kwargs = {"device": device, "dtype": dtype} 350 | super().__init__() 351 | self.embed_dim = embed_dim 352 | self.num_heads = num_heads 353 | self.causal = causal 354 | 355 | assert ( 356 | self.embed_dim % num_heads == 0 357 | ), "self.kdim must be divisible by num_heads" 358 | self.head_dim = self.embed_dim // num_heads 359 | assert self.head_dim in [16, 32, 64], "Only support head_dim == 16, 32, or 64" 360 | 361 | # self.inner_attn = FlashAttention( 362 | # attention_dropout=attention_dropout 363 | # ) 364 | self.attention_dropout = attention_dropout 365 | self.inner_attn = flash_attn_qkvpacked_func 366 | self.rearrange = rearrange 367 | 368 | def forward(self, qkv, attn_mask=None, key_padding_mask=None, need_weights=False): 369 | qkv = self.rearrange( 370 | qkv, "b (three h d) s -> b s three h d", three=3, h=self.num_heads 371 | ).contiguous() 372 | # qkv, _ = self.inner_attn( 373 | # qkv, 374 | # key_padding_mask=key_padding_mask, 375 | # need_weights=need_weights, 376 | # causal=self.causal, 377 | # ) 378 | attention_dropout = 0.0 379 | if self.training: 380 | attention_dropout = self.attention_dropout 381 | qkv = self.inner_attn(qkv, attention_dropout) 382 | return self.rearrange(qkv, "b s h d -> b (h d) s") 383 | 384 | 385 | def count_flops_attn(model, _x, y): 386 | """ 387 | A counter for the `thop` package to count the operations in an 388 | attention operation. 389 | Meant to be used like: 390 | macs, params = thop.profile( 391 | model, 392 | inputs=(inputs, timestamps), 393 | custom_ops={QKVAttention: QKVAttention.count_flops}, 394 | ) 395 | """ 396 | b, c, *spatial = y[0].shape 397 | num_spatial = int(np.prod(spatial)) 398 | # We perform two matmuls with the same number of ops. 399 | # The first computes the weight matrix, the second computes 400 | # the combination of the value vectors. 401 | matmul_ops = 2 * b * (num_spatial**2) * c 402 | model.total_ops += th.DoubleTensor([matmul_ops]) 403 | 404 | 405 | class QKVAttentionLegacy(nn.Module): 406 | """ 407 | A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping 408 | """ 409 | 410 | def __init__(self, n_heads): 411 | super().__init__() 412 | from einops import rearrange 413 | self.n_heads = n_heads 414 | self.rearrange = rearrange 415 | 416 | def forward(self, qkv): 417 | """ 418 | Apply QKV attention. 419 | 420 | :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs. 421 | :return: an [N x (H * C) x T] tensor after attention. 422 | """ 423 | bs, width, length = qkv.shape 424 | assert width % (3 * self.n_heads) == 0 425 | ch = width // (3 * self.n_heads) 426 | q, k, v = self.rearrange( 427 | qkv, "b (three h d) s -> (b h) (three d) s", three=3, h=self.n_heads 428 | ).contiguous().split(ch, dim=1) 429 | 430 | # q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1) 431 | scale = 1 / math.sqrt(math.sqrt(ch)) 432 | weight = th.einsum( 433 | "bct,bcs->bts", q * scale, k * scale 434 | ) # More stable with f16 than dividing afterwards 435 | weight = th.softmax(weight.float(), dim=-1).type(weight.dtype) 436 | a = th.einsum("bts,bcs->bct", weight, v) 437 | return a.reshape(bs, -1, length) 438 | 439 | @staticmethod 440 | def count_flops(model, _x, y): 441 | return count_flops_attn(model, _x, y) 442 | 443 | 444 | # class QKVAttention(nn.Module): 445 | # """ 446 | # A module which performs QKV attention and splits in a different order. 447 | # """ 448 | 449 | # def __init__(self, n_heads): 450 | # super().__init__() 451 | # self.n_heads = n_heads 452 | 453 | # def forward(self, qkv): 454 | # """ 455 | # Apply QKV attention. 456 | 457 | # :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs. 458 | # :return: an [N x (H * C) x T] tensor after attention. 459 | # """ 460 | # bs, width, length = qkv.shape 461 | # assert width % (3 * self.n_heads) == 0 462 | # ch = width // (3 * self.n_heads) 463 | # q, k, v = qkv.chunk(3, dim=1) 464 | # scale = 1 / math.sqrt(math.sqrt(ch)) 465 | # weight = th.einsum( 466 | # "bct,bcs->bts", 467 | # (q * scale).view(bs * self.n_heads, ch, length), 468 | # (k * scale).view(bs * self.n_heads, ch, length), 469 | # ) # More stable with f16 than dividing afterwards 470 | # weight = th.softmax(weight.float(), dim=-1).type(weight.dtype) 471 | # a = th.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length)) 472 | # return a.reshape(bs, -1, length) 473 | 474 | # @staticmethod 475 | # def count_flops(model, _x, y): 476 | # return count_flops_attn(model, _x, y) 477 | 478 | 479 | class QKVAttention(nn.Module): 480 | """ 481 | A module which performs QKV attention. Fallback from Blocksparse if use_fp16=False 482 | """ 483 | 484 | def __init__(self, n_heads): 485 | super().__init__() 486 | self.n_heads = n_heads 487 | 488 | def forward(self, qkv, encoder_kv=None): 489 | """ 490 | Apply QKV attention. 491 | 492 | :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs. 493 | :return: an [N x (H * C) x T] tensor after attention. 494 | """ 495 | bs, width, length = qkv.shape 496 | assert width % (3 * self.n_heads) == 0 497 | ch = width // (3 * self.n_heads) 498 | q, k, v = qkv.chunk(3, dim=1) 499 | if encoder_kv is not None: 500 | assert encoder_kv.shape[1] == 2 * ch * self.n_heads 501 | ek, ev = encoder_kv.chunk(2, dim=1) 502 | k = th.cat([ek, k], dim=-1) 503 | v = th.cat([ev, v], dim=-1) 504 | scale = 1 / math.sqrt(math.sqrt(ch)) 505 | weight = th.einsum( 506 | "bct,bcs->bts", 507 | (q * scale).view(bs * self.n_heads, ch, length), 508 | (k * scale).view(bs * self.n_heads, ch, -1), 509 | ) # More stable with f16 than dividing afterwards 510 | weight = th.softmax(weight.float(), dim=-1).type(weight.dtype) 511 | a = th.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, -1)) 512 | return a.reshape(bs, -1, length) 513 | 514 | @staticmethod 515 | def count_flops(model, _x, y): 516 | return count_flops_attn(model, _x, y) 517 | 518 | 519 | class UNetModel(nn.Module): 520 | """ 521 | The full UNet model with attention and timestep embedding. 522 | 523 | :param in_channels: channels in the input Tensor. 524 | :param model_channels: base channel count for the model. 525 | :param out_channels: channels in the output Tensor. 526 | :param num_res_blocks: number of residual blocks per downsample. 527 | :param attention_resolutions: a collection of downsample rates at which 528 | attention will take place. May be a set, list, or tuple. 529 | For example, if this contains 4, then at 4x downsampling, attention 530 | will be used. 531 | :param dropout: the dropout probability. 532 | :param channel_mult: channel multiplier for each level of the UNet. 533 | :param conv_resample: if True, use learned convolutions for upsampling and 534 | downsampling. 535 | :param dims: determines if the signal is 1D, 2D, or 3D. 536 | :param num_classes: if specified (as an int), then this model will be 537 | class-conditional with `num_classes` classes. 538 | :param use_checkpoint: use gradient checkpointing to reduce memory usage. 539 | :param num_heads: the number of attention heads in each attention layer. 540 | :param num_heads_channels: if specified, ignore num_heads and instead use 541 | a fixed channel width per attention head. 542 | :param num_heads_upsample: works with num_heads to set a different number 543 | of heads for upsampling. Deprecated. 544 | :param use_scale_shift_norm: use a FiLM-like conditioning mechanism. 545 | :param resblock_updown: use residual blocks for up/downsampling. 546 | :param use_new_attention_order: use a different attention pattern for potentially 547 | increased efficiency. 548 | """ 549 | 550 | def __init__( 551 | self, 552 | image_size, 553 | in_channels, 554 | model_channels, 555 | out_channels, 556 | num_res_blocks, 557 | attention_resolutions, 558 | dropout=0, 559 | channel_mult=(1, 2, 4, 8), 560 | conv_resample=True, 561 | dims=2, 562 | num_classes=None, 563 | use_checkpoint=False, 564 | use_fp16=False, 565 | num_heads=1, 566 | num_head_channels=-1, 567 | num_heads_upsample=-1, 568 | use_scale_shift_norm=False, 569 | resblock_updown=False, 570 | use_new_attention_order=False, 571 | ): 572 | super().__init__() 573 | 574 | if num_heads_upsample == -1: 575 | num_heads_upsample = num_heads 576 | 577 | self.image_size = image_size 578 | self.in_channels = in_channels 579 | self.model_channels = model_channels 580 | self.out_channels = out_channels 581 | self.num_res_blocks = num_res_blocks 582 | self.attention_resolutions = attention_resolutions 583 | self.dropout = dropout 584 | self.channel_mult = channel_mult 585 | self.conv_resample = conv_resample 586 | self.num_classes = num_classes 587 | self.use_checkpoint = use_checkpoint 588 | self.dtype = th.float16 if use_fp16 else th.float32 589 | self.num_heads = num_heads 590 | self.num_head_channels = num_head_channels 591 | self.num_heads_upsample = num_heads_upsample 592 | 593 | time_embed_dim = model_channels * 4 594 | self.time_embed = nn.Sequential( 595 | linear(model_channels, time_embed_dim), 596 | nn.SiLU(), 597 | linear(time_embed_dim, time_embed_dim), 598 | ) 599 | 600 | if self.num_classes is not None: 601 | self.label_emb = nn.Embedding(num_classes, time_embed_dim) 602 | 603 | ch = input_ch = int(channel_mult[0] * model_channels) 604 | self.input_blocks = nn.ModuleList( 605 | [TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))] 606 | ) 607 | self._feature_size = ch 608 | input_block_chans = [ch] 609 | ds = 1 610 | for level, mult in enumerate(channel_mult): 611 | for _ in range(num_res_blocks): 612 | layers = [ 613 | ResBlock( 614 | ch, 615 | time_embed_dim, 616 | dropout, 617 | out_channels=int(mult * model_channels), 618 | dims=dims, 619 | use_checkpoint=use_checkpoint, 620 | use_scale_shift_norm=use_scale_shift_norm, 621 | ) 622 | ] 623 | ch = int(mult * model_channels) 624 | if ds in attention_resolutions: 625 | layers.append( 626 | AttentionBlock( 627 | ch, 628 | use_checkpoint=use_checkpoint, 629 | num_heads=num_heads, 630 | num_head_channels=num_head_channels, 631 | use_new_attention_order=use_new_attention_order, 632 | ) 633 | ) 634 | self.input_blocks.append(TimestepEmbedSequential(*layers)) 635 | self._feature_size += ch 636 | input_block_chans.append(ch) 637 | if level != len(channel_mult) - 1: 638 | out_ch = ch 639 | self.input_blocks.append( 640 | TimestepEmbedSequential( 641 | ResBlock( 642 | ch, 643 | time_embed_dim, 644 | dropout, 645 | out_channels=out_ch, 646 | dims=dims, 647 | use_checkpoint=use_checkpoint, 648 | use_scale_shift_norm=use_scale_shift_norm, 649 | down=True, 650 | ) 651 | if resblock_updown 652 | else Downsample( 653 | ch, conv_resample, dims=dims, out_channels=out_ch 654 | ) 655 | ) 656 | ) 657 | ch = out_ch 658 | input_block_chans.append(ch) 659 | ds *= 2 660 | self._feature_size += ch 661 | 662 | self.middle_block = TimestepEmbedSequential( 663 | ResBlock( 664 | ch, 665 | time_embed_dim, 666 | dropout, 667 | dims=dims, 668 | use_checkpoint=use_checkpoint, 669 | use_scale_shift_norm=use_scale_shift_norm, 670 | ), 671 | AttentionBlock( 672 | ch, 673 | use_checkpoint=use_checkpoint, 674 | num_heads=num_heads, 675 | num_head_channels=num_head_channels, 676 | use_new_attention_order=use_new_attention_order, 677 | ), 678 | ResBlock( 679 | ch, 680 | time_embed_dim, 681 | dropout, 682 | dims=dims, 683 | use_checkpoint=use_checkpoint, 684 | use_scale_shift_norm=use_scale_shift_norm, 685 | ), 686 | ) 687 | self._feature_size += ch 688 | 689 | self.output_blocks = nn.ModuleList([]) 690 | for level, mult in list(enumerate(channel_mult))[::-1]: 691 | for i in range(num_res_blocks + 1): 692 | ich = input_block_chans.pop() 693 | layers = [ 694 | ResBlock( 695 | ch + ich, 696 | time_embed_dim, 697 | dropout, 698 | out_channels=int(model_channels * mult), 699 | dims=dims, 700 | use_checkpoint=use_checkpoint, 701 | use_scale_shift_norm=use_scale_shift_norm, 702 | ) 703 | ] 704 | ch = int(model_channels * mult) 705 | if ds in attention_resolutions: 706 | layers.append( 707 | AttentionBlock( 708 | ch, 709 | use_checkpoint=use_checkpoint, 710 | num_heads=num_heads_upsample, 711 | num_head_channels=num_head_channels, 712 | use_new_attention_order=use_new_attention_order, 713 | ) 714 | ) 715 | if level and i == num_res_blocks: 716 | out_ch = ch 717 | layers.append( 718 | ResBlock( 719 | ch, 720 | time_embed_dim, 721 | dropout, 722 | out_channels=out_ch, 723 | dims=dims, 724 | use_checkpoint=use_checkpoint, 725 | use_scale_shift_norm=use_scale_shift_norm, 726 | up=True, 727 | ) 728 | if resblock_updown 729 | else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch) 730 | ) 731 | ds //= 2 732 | self.output_blocks.append(TimestepEmbedSequential(*layers)) 733 | self._feature_size += ch 734 | 735 | self.out = nn.Sequential( 736 | normalization(ch), 737 | nn.SiLU(), 738 | zero_module(conv_nd(dims, input_ch, out_channels, 3, padding=1)), 739 | ) 740 | 741 | def convert_to_fp16(self): 742 | """ 743 | Convert the torso of the model to float16. 744 | """ 745 | self.input_blocks.apply(convert_module_to_f16) 746 | self.middle_block.apply(convert_module_to_f16) 747 | self.output_blocks.apply(convert_module_to_f16) 748 | 749 | def convert_to_fp32(self): 750 | """ 751 | Convert the torso of the model to float32. 752 | """ 753 | self.input_blocks.apply(convert_module_to_f32) 754 | self.middle_block.apply(convert_module_to_f32) 755 | self.output_blocks.apply(convert_module_to_f32) 756 | 757 | def forward(self, x, timesteps, y=None): 758 | """ 759 | Apply the model to an input batch. 760 | 761 | :param x: an [N x C x ...] Tensor of inputs. 762 | :param timesteps: a 1-D batch of timesteps. 763 | :param y: an [N] Tensor of labels, if class-conditional. 764 | :return: an [N x C x ...] Tensor of outputs. 765 | """ 766 | assert (y is not None) == ( 767 | self.num_classes is not None 768 | ), "must specify y if and only if the model is class-conditional" 769 | 770 | hs = [] 771 | emb = self.time_embed(timestep_embedding(timesteps, self.model_channels)) 772 | 773 | if self.num_classes is not None: 774 | assert y.shape == (x.shape[0],) 775 | emb = emb + self.label_emb(y) 776 | 777 | h = x.type(self.dtype) 778 | for module in self.input_blocks: 779 | h = module(h, emb) 780 | hs.append(h) 781 | h = self.middle_block(h, emb) 782 | for module in self.output_blocks: 783 | h = th.cat([h, hs.pop()], dim=1) 784 | h = module(h, emb) 785 | h = h.type(x.dtype) 786 | return self.out(h) 787 | -------------------------------------------------------------------------------- /cm/network.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 | # 3 | # This work is licensed under a Creative Commons 4 | # Attribution-NonCommercial-ShareAlike 4.0 International License. 5 | # You should have received a copy of the license along with this 6 | # work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/ 7 | 8 | """Model architectures and preconditioning schemes used in the paper 9 | "Elucidating the Design Space of Diffusion-Based Generative Models".""" 10 | 11 | import numpy as np 12 | import torch 13 | # from torch_utils import persistence 14 | from torch.nn.functional import silu 15 | import torch.nn as nn 16 | # from .fp16_util import convert_module_to_f16, convert_module_to_f32 17 | # from .nn_test import ( 18 | # checkpoint, 19 | # conv_nd, 20 | # linear, 21 | # avg_pool_nd, 22 | # zero_module, 23 | # normalization, 24 | # timestep_embedding, 25 | # ) 26 | 27 | #---------------------------------------------------------------------------- 28 | # Unified routine for initializing weights and biases. 29 | 30 | def weight_init(shape, mode, fan_in, fan_out): 31 | if mode == 'xavier_uniform': return np.sqrt(6 / (fan_in + fan_out)) * (torch.rand(*shape) * 2 - 1) 32 | if mode == 'xavier_normal': return np.sqrt(2 / (fan_in + fan_out)) * torch.randn(*shape) 33 | if mode == 'kaiming_uniform': return np.sqrt(3 / fan_in) * (torch.rand(*shape) * 2 - 1) 34 | if mode == 'kaiming_normal': return np.sqrt(1 / fan_in) * torch.randn(*shape) 35 | raise ValueError(f'Invalid init mode "{mode}"') 36 | 37 | 38 | def convert_module_to_f16(l): 39 | """ 40 | Convert primitive modules to float16. 41 | """ 42 | if isinstance(l, (Conv2d)): 43 | if l.weight is not None: 44 | l.weight.data = l.weight.data.half() 45 | if l.bias is not None: 46 | l.bias.data = l.bias.data.half() 47 | 48 | 49 | def convert_module_to_f32(l): 50 | """ 51 | Convert primitive modules to float32, undoing convert_module_to_f16(). 52 | """ 53 | if isinstance(l, (Conv2d)): 54 | if l.weight is not None: 55 | l.weight.data = l.weight.data.float() 56 | if l.bias is not None: 57 | l.bias.data = l.bias.data.float() 58 | 59 | #---------------------------------------------------------------------------- 60 | # Fully-connected layer. 61 | 62 | # @persistence.persistent_class 63 | class Linear(torch.nn.Module): 64 | def __init__(self, in_features, out_features, bias=True, init_mode='kaiming_normal', init_weight=1, init_bias=0): 65 | super().__init__() 66 | self.in_features = in_features 67 | self.out_features = out_features 68 | init_kwargs = dict(mode=init_mode, fan_in=in_features, fan_out=out_features) 69 | self.weight = torch.nn.Parameter(weight_init([out_features, in_features], **init_kwargs) * init_weight) 70 | self.bias = torch.nn.Parameter(weight_init([out_features], **init_kwargs) * init_bias) if bias else None 71 | 72 | def forward(self, x): 73 | x = x @ self.weight.t() 74 | if self.bias is not None: 75 | x = x.add_(self.bias) 76 | return x 77 | 78 | #---------------------------------------------------------------------------- 79 | # Convolutional layer with optional up/downsampling. 80 | 81 | # @persistence.persistent_class 82 | class Conv2d(torch.nn.Module): 83 | def __init__(self, 84 | in_channels, out_channels, kernel, bias=True, up=False, down=False, 85 | resample_filter=[1,1], fused_resample=False, init_mode='kaiming_normal', init_weight=1, init_bias=0, 86 | ): 87 | assert not (up and down) 88 | super().__init__() 89 | self.in_channels = in_channels 90 | self.out_channels = out_channels 91 | self.up = up 92 | self.down = down 93 | self.fused_resample = fused_resample 94 | init_kwargs = dict(mode=init_mode, fan_in=in_channels*kernel*kernel, fan_out=out_channels*kernel*kernel) 95 | self.weight = torch.nn.Parameter(weight_init([out_channels, in_channels, kernel, kernel], **init_kwargs) * init_weight) if kernel else None 96 | self.bias = torch.nn.Parameter(weight_init([out_channels], **init_kwargs) * init_bias) if kernel and bias else None 97 | f = torch.as_tensor(resample_filter, dtype=torch.float32) 98 | f = f.ger(f).unsqueeze(0).unsqueeze(1) / f.sum().square() 99 | self.register_buffer('resample_filter', f if up or down else None) 100 | 101 | def forward(self, x): 102 | w = self.weight.to(x.dtype) if self.weight is not None else None 103 | b = self.bias.to(x.dtype) if self.bias is not None else None 104 | f = self.resample_filter.to(x.dtype) if self.resample_filter is not None else None 105 | w_pad = w.shape[-1] // 2 if w is not None else 0 106 | f_pad = (f.shape[-1] - 1) // 2 if f is not None else 0 107 | 108 | if self.fused_resample and self.up and w is not None: 109 | x = torch.nn.functional.conv_transpose2d(x, f.mul(4).tile([self.in_channels, 1, 1, 1]), groups=self.in_channels, stride=2, padding=max(f_pad - w_pad, 0)) 110 | x = torch.nn.functional.conv2d(x, w, padding=max(w_pad - f_pad, 0)) 111 | elif self.fused_resample and self.down and w is not None: 112 | x = torch.nn.functional.conv2d(x, w, padding=w_pad+f_pad) 113 | x = torch.nn.functional.conv2d(x, f.tile([self.out_channels, 1, 1, 1]), groups=self.out_channels, stride=2) 114 | else: 115 | if self.up: 116 | x = torch.nn.functional.conv_transpose2d(x, f.mul(4).tile([self.in_channels, 1, 1, 1]), groups=self.in_channels, stride=2, padding=f_pad) 117 | if self.down: 118 | x = torch.nn.functional.conv2d(x, f.tile([self.in_channels, 1, 1, 1]), groups=self.in_channels, stride=2, padding=f_pad) 119 | if w is not None: 120 | x = torch.nn.functional.conv2d(x, w, padding=w_pad) 121 | if b is not None: 122 | x = x.add_(b.reshape(1, -1, 1, 1)) 123 | return x 124 | 125 | #---------------------------------------------------------------------------- 126 | # Group normalization. 127 | 128 | # @persistence.persistent_class 129 | class GroupNorm(torch.nn.Module): 130 | def __init__(self, num_channels, num_groups=32, min_channels_per_group=4, eps=1e-5): 131 | super().__init__() 132 | self.num_groups = min(num_groups, num_channels // min_channels_per_group) 133 | self.eps = eps 134 | self.weight = torch.nn.Parameter(torch.ones(num_channels)) 135 | self.bias = torch.nn.Parameter(torch.zeros(num_channels)) 136 | 137 | def forward(self, x): 138 | x = torch.nn.functional.group_norm(x, num_groups=self.num_groups, weight=self.weight.to(x.dtype), bias=self.bias.to(x.dtype), eps=self.eps) 139 | return x 140 | 141 | #---------------------------------------------------------------------------- 142 | # Attention weight computation, i.e., softmax(Q^T * K). 143 | # Performs all computation using FP32, but uses the original datatype for 144 | # inputs/outputs/gradients to conserve memory. 145 | 146 | class AttentionOp(torch.autograd.Function): 147 | @staticmethod 148 | def forward(ctx, q, k): 149 | # w = torch.einsum('ncq,nck->nqk', q.to(torch.float32), (k / np.sqrt(k.shape[1])).to(torch.float32)).softmax(dim=2).to(q.dtype) 150 | 151 | w = torch.einsum('ncq,nck->nqk', q.to(torch.float32), (k / np.sqrt(k.shape[1])).to(torch.float32)) 152 | # w = w.softmax(dim=2) 153 | w = our_softmax(w, dim=2) 154 | w = w.to(q.dtype) 155 | ctx.save_for_backward(q, k, w) 156 | return w 157 | 158 | @staticmethod 159 | def backward(ctx, dw): 160 | q, k, w = ctx.saved_tensors 161 | db = torch._softmax_backward_data(grad_output=dw.to(torch.float32), output=w.to(torch.float32), dim=2, input_dtype=torch.float32) 162 | dq = torch.einsum('nck,nqk->ncq', k.to(torch.float32), db).to(q.dtype) / np.sqrt(k.shape[1]) 163 | dk = torch.einsum('ncq,nqk->nck', q.to(torch.float32), db).to(k.dtype) / np.sqrt(k.shape[1]) 164 | return dq, dk 165 | 166 | @torch.jit.script 167 | def our_softmax(x, dim:int=-1): 168 | """ 169 | x: (B, C, C) 170 | """ 171 | 172 | maxes = torch.max(x, dim, keepdim=True)[0] 173 | x_exp = torch.exp(x-maxes) 174 | x_exp_sum = torch.sum(x_exp, dim, keepdim=True) 175 | output_custom = x_exp/x_exp_sum 176 | return output_custom 177 | 178 | #---------------------------------------------------------------------------- 179 | # Unified U-Net block with optional up/downsampling and self-attention. 180 | # Represents the union of all features employed by the DDPM++, NCSN++, and 181 | # ADM architectures. 182 | 183 | # @persistence.persistent_class 184 | class UNetBlock(torch.nn.Module): 185 | def __init__(self, 186 | in_channels, out_channels, emb_channels, up=False, down=False, attention=False, 187 | num_heads=None, channels_per_head=64, dropout=0, skip_scale=1, eps=1e-5, 188 | resample_filter=[1,1], resample_proj=False, adaptive_scale=True, 189 | init=dict(), init_zero=dict(init_weight=0), init_attn=None, 190 | # flash_atten=True, 191 | flash_atten=False, 192 | ): 193 | super().__init__() 194 | self.in_channels = in_channels 195 | self.out_channels = out_channels 196 | self.emb_channels = emb_channels 197 | self.num_heads = 0 if not attention else num_heads if num_heads is not None else out_channels // channels_per_head 198 | self.dropout = dropout 199 | self.skip_scale = skip_scale 200 | self.adaptive_scale = adaptive_scale 201 | self.flash_atten = flash_atten 202 | if self.flash_atten and self.num_heads: 203 | from einops import rearrange 204 | import sys 205 | sys.path.insert(0, "/home/work/workspace/flash-attention") 206 | from flash_attn import flash_attn_varlen_qkvpacked_func, flash_attn_qkvpacked_func 207 | 208 | self.rearrange = rearrange 209 | # self.inner_attn = flash_attn_varlen_qkvpacked_func 210 | self.inner_attn = flash_attn_qkvpacked_func 211 | 212 | self.norm0 = GroupNorm(num_channels=in_channels, eps=eps) 213 | self.conv0 = Conv2d(in_channels=in_channels, out_channels=out_channels, kernel=3, up=up, down=down, resample_filter=resample_filter, **init) 214 | self.affine = Linear(in_features=emb_channels, out_features=out_channels*(2 if adaptive_scale else 1), **init) 215 | self.norm1 = GroupNorm(num_channels=out_channels, eps=eps) 216 | self.conv1 = Conv2d(in_channels=out_channels, out_channels=out_channels, kernel=3, **init_zero) 217 | 218 | self.skip = None 219 | if out_channels != in_channels or up or down: 220 | kernel = 1 if resample_proj or out_channels!= in_channels else 0 221 | self.skip = Conv2d(in_channels=in_channels, out_channels=out_channels, kernel=kernel, up=up, down=down, resample_filter=resample_filter, **init) 222 | 223 | if self.num_heads: 224 | self.norm2 = GroupNorm(num_channels=out_channels, eps=eps) 225 | self.qkv = Conv2d(in_channels=out_channels, out_channels=out_channels*3, kernel=1, **(init_attn if init_attn is not None else init)) 226 | self.proj = Conv2d(in_channels=out_channels, out_channels=out_channels, kernel=1, **init_zero) 227 | 228 | def forward(self, x, emb): 229 | orig = x 230 | x = self.conv0(silu(self.norm0(x))) 231 | 232 | params = self.affine(emb).unsqueeze(2).unsqueeze(3).to(x.dtype) 233 | if self.adaptive_scale: 234 | scale, shift = params.chunk(chunks=2, dim=1) 235 | x = silu(torch.addcmul(shift, self.norm1(x), scale + 1)) 236 | else: 237 | x = silu(self.norm1(x.add_(params))) 238 | 239 | x = self.conv1(torch.nn.functional.dropout(x, p=self.dropout, training=self.training)) 240 | x = x.add_(self.skip(orig) if self.skip is not None else orig) 241 | x = x * self.skip_scale 242 | 243 | if self.num_heads: 244 | if self.flash_atten: 245 | b, _, *spatial = x.shape 246 | qkv = self.qkv(self.norm2(x)).view(b, -1, np.prod(spatial)) 247 | qkv = self.rearrange( 248 | qkv, "b (h d three) s -> b s three h d", three=3, h=self.num_heads 249 | ) 250 | a = self.inner_attn(qkv) 251 | a = self.rearrange(a, "b s h d -> b (h d) s") 252 | a = a.view(b, -1, *spatial) 253 | 254 | else: 255 | q, k, v = self.qkv(self.norm2(x)).reshape(x.shape[0] * self.num_heads, x.shape[1] // self.num_heads, 3, -1).unbind(2) 256 | # w = AttentionOp.apply(q, k) 257 | 258 | w = torch.einsum('ncq,nck->nqk', q.to(torch.float32), (k / np.sqrt(k.shape[1])).to(torch.float32)) 259 | w = our_softmax(w, dim=2) 260 | w = w.to(q.dtype) 261 | 262 | a = torch.einsum('nqk,nck->ncq', w, v) 263 | a = a.reshape(*x.shape).contiguous() 264 | # q, k, v = self.qkv(self.norm2(x)).reshape(x.shape[0] * self.num_heads, x.shape[1] // self.num_heads, 3, -1).unbind(2) 265 | # w = AttentionOp.apply(q, k) 266 | # a = torch.einsum('nqk,nck->ncq', w, v) 267 | # x = self.proj(a.reshape(*x.shape)).add_(x) 268 | x = self.proj(a).add_(x) 269 | x = x * self.skip_scale 270 | 271 | # if x.isnan().sum()>0: 272 | # import ipdb; ipdb.set_trace() 273 | return x 274 | 275 | #---------------------------------------------------------------------------- 276 | # Timestep embedding used in the DDPM++ and ADM architectures. 277 | 278 | # @persistence.persistent_class 279 | class PositionalEmbedding(torch.nn.Module): 280 | def __init__(self, num_channels, max_positions=10000, endpoint=False): 281 | super().__init__() 282 | self.num_channels = num_channels 283 | self.max_positions = max_positions 284 | self.endpoint = endpoint 285 | 286 | def forward(self, x): 287 | freqs = torch.arange(start=0, end=self.num_channels//2, dtype=torch.float32, device=x.device) 288 | freqs = freqs / (self.num_channels // 2 - (1 if self.endpoint else 0)) 289 | freqs = (1 / self.max_positions) ** freqs 290 | x = x.ger(freqs) 291 | x = torch.cat([x.cos(), x.sin()], dim=1) 292 | return x 293 | 294 | #---------------------------------------------------------------------------- 295 | # Timestep embedding used in the NCSN++ architecture. 296 | 297 | # @persistence.persistent_class 298 | class FourierEmbedding(torch.nn.Module): 299 | def __init__(self, num_channels, scale=16): 300 | super().__init__() 301 | self.register_buffer('freqs', torch.randn(num_channels // 2) * scale) 302 | 303 | def forward(self, x): 304 | x = x.ger((2 * np.pi * self.freqs)) 305 | x = torch.cat([x.cos(), x.sin()], dim=1) 306 | return x 307 | 308 | #---------------------------------------------------------------------------- 309 | # Reimplementation of the DDPM++ and NCSN++ architectures from the paper 310 | # "Score-Based Generative Modeling through Stochastic Differential 311 | # Equations". Equivalent to the original implementation by Song et al., 312 | # available at https://github.com/yang-song/score_sde_pytorch 313 | 314 | #---------------------------------------------------------------------------- 315 | # Reimplementation of the ADM architecture from the paper 316 | # "Diffusion Models Beat GANS on Image Synthesis". Equivalent to the 317 | # original implementation by Dhariwal and Nichol, available at 318 | # https://github.com/openai/guided-diffusion 319 | 320 | # @persistence.persistent_class 321 | class DhariwalUNet(torch.nn.Module): 322 | def __init__(self, 323 | image_size, 324 | in_channels, 325 | model_channels, 326 | out_channels, 327 | num_res_blocks, 328 | attention_resolutions, 329 | dropout=0, 330 | channel_mult=(1, 2, 4, 8), 331 | conv_resample=True, 332 | dims=2, 333 | num_classes=None, 334 | use_checkpoint=False, 335 | use_fp16=False, 336 | num_heads=1, 337 | num_head_channels=-1, 338 | num_heads_upsample=-1, 339 | use_scale_shift_norm=False, 340 | resblock_updown=False, 341 | use_new_attention_order=False, 342 | ): 343 | 344 | # img_resolution, # Image resolution at input/output. 345 | # in_channels, # Number of color channels at input. 346 | # out_channels, # Number of color channels at output. 347 | # label_dim = 0, # Number of class labels, 0 = unconditional. 348 | # augment_dim = 0, # Augmentation label dimensionality, 0 = no augmentation. 349 | 350 | # model_channels = 192, # Base multiplier for the number of channels. 351 | # channel_mult = [1,2,3,4], # Per-resolution multipliers for the number of channels. 352 | # channel_mult_emb = 4, # Multiplier for the dimensionality of the embedding vector. 353 | # num_blocks = 3, # Number of residual blocks per resolution. 354 | # attn_resolutions = [32,16,8], # List of resolutions with self-attention. 355 | # dropout = 0.10, # List of resolutions with self-attention. 356 | # label_dropout = 0, # Dropout probability of class labels for classifier-free guidance. 357 | 358 | super().__init__() 359 | self.label_dropout = 0 360 | self.model_channels = model_channels 361 | self.channel_mult_emb = 4 362 | self.augment_dim =0 363 | self.in_channels = in_channels 364 | self.image_size = image_size 365 | self.num_classes = num_classes 366 | self.dtype = torch.float16 if use_fp16 else torch.float32 367 | 368 | attn_resolutions = [32,16,8] 369 | img_resolution = image_size 370 | num_blocks = num_res_blocks 371 | 372 | 373 | emb_channels = model_channels * self.channel_mult_emb 374 | init = dict(init_mode='kaiming_uniform', init_weight=np.sqrt(1/3), init_bias=np.sqrt(1/3)) 375 | init_zero = dict(init_mode='kaiming_uniform', init_weight=0, init_bias=0) 376 | block_kwargs = dict(emb_channels=emb_channels, channels_per_head=64, dropout=dropout, init=init, init_zero=init_zero) 377 | 378 | # Mapping. 379 | self.map_noise = PositionalEmbedding(num_channels=model_channels) 380 | # time_embed_dim = model_channels * 4 381 | # self.time_embed = nn.Sequential( 382 | # linear(model_channels, time_embed_dim), 383 | # nn.SiLU(), 384 | # linear(time_embed_dim, time_embed_dim), 385 | # ) 386 | 387 | 388 | self.map_augment = Linear(in_features=self.augment_dim, out_features=model_channels, bias=False, **init_zero) if self.augment_dim else None 389 | 390 | self.map_layer0 = Linear(in_features=model_channels, out_features=emb_channels, **init) 391 | 392 | 393 | self.map_layer1 = Linear(in_features=emb_channels, out_features=emb_channels, **init) 394 | 395 | 396 | self.map_label = Linear(in_features=num_classes, out_features=emb_channels, bias=False, init_mode='kaiming_normal', init_weight=np.sqrt(num_classes)) if num_classes else None 397 | 398 | # Encoder. 399 | self.enc = torch.nn.ModuleDict() 400 | cout = in_channels 401 | for level, mult in enumerate(channel_mult): 402 | res = img_resolution >> level 403 | if level == 0: 404 | cin = cout 405 | cout = model_channels * mult 406 | self.enc[f'{res}x{res}_conv'] = Conv2d(in_channels=cin, out_channels=cout, kernel=3, **init) 407 | else: 408 | self.enc[f'{res}x{res}_down'] = UNetBlock(in_channels=cout, out_channels=cout, down=True, **block_kwargs) 409 | for idx in range(num_blocks): 410 | cin = cout 411 | cout = model_channels * mult 412 | self.enc[f'{res}x{res}_block{idx}'] = UNetBlock(in_channels=cin, out_channels=cout, attention=(res in attn_resolutions), **block_kwargs) 413 | skips = [block.out_channels for block in self.enc.values()] 414 | 415 | # Decoder. 416 | self.dec = torch.nn.ModuleDict() 417 | for level, mult in reversed(list(enumerate(channel_mult))): 418 | res = img_resolution >> level 419 | if level == len(channel_mult) - 1: 420 | self.dec[f'{res}x{res}_in0'] = UNetBlock(in_channels=cout, out_channels=cout, attention=True, **block_kwargs) 421 | self.dec[f'{res}x{res}_in1'] = UNetBlock(in_channels=cout, out_channels=cout, **block_kwargs) 422 | else: 423 | self.dec[f'{res}x{res}_up'] = UNetBlock(in_channels=cout, out_channels=cout, up=True, **block_kwargs) 424 | for idx in range(num_blocks + 1): 425 | cin = cout + skips.pop() 426 | cout = model_channels * mult 427 | self.dec[f'{res}x{res}_block{idx}'] = UNetBlock(in_channels=cin, out_channels=cout, attention=(res in attn_resolutions), **block_kwargs) 428 | self.out_norm = GroupNorm(num_channels=cout) 429 | self.out_conv = Conv2d(in_channels=cout, out_channels=out_channels, kernel=3, **init_zero) 430 | 431 | def convert_to_fp16(self): 432 | """ 433 | Convert the torso of the model to float16. 434 | """ 435 | self.enc.apply(convert_module_to_f16) 436 | # self.middle_block.apply(convert_module_to_f16) 437 | self.dec.apply(convert_module_to_f16) 438 | 439 | def convert_to_fp32(self): 440 | """ 441 | Convert the torso of the model to float32. 442 | """ 443 | self.enc.apply(convert_module_to_f32) 444 | # self.middle_block.apply(convert_module_to_f32) 445 | self.dec.apply(convert_module_to_f32) 446 | 447 | def forward(self, x, timesteps, y): 448 | # Mapping. 449 | class_labels = torch.nn.functional.one_hot(y, num_classes=self.num_classes).to(x.dtype) 450 | noise_labels = timesteps 451 | 452 | emb = self.map_noise(noise_labels) 453 | # emb = self.time_embed(timestep_embedding(timesteps, self.model_channels)) 454 | 455 | # if self.map_augment is not None and augment_labels is not None: 456 | # emb = emb + self.map_augment(augment_labels) 457 | emb = silu(self.map_layer0(emb)) 458 | emb = self.map_layer1(emb) 459 | 460 | 461 | if self.map_label is not None: 462 | tmp = class_labels 463 | if self.training and self.label_dropout: 464 | tmp = tmp * (torch.rand([x.shape[0], 1], device=x.device) >= self.label_dropout) 465 | emb = emb + self.map_label(tmp) 466 | emb = silu(emb) 467 | 468 | # Encoder. 469 | skips = [] 470 | x = x.type(self.dtype) 471 | for block in self.enc.values(): 472 | x = block(x, emb) if isinstance(block, UNetBlock) else block(x) 473 | skips.append(x) 474 | 475 | # Decoder. 476 | for block in self.dec.values(): 477 | if x.shape[1] != block.in_channels: 478 | x = torch.cat([x, skips.pop()], dim=1) 479 | x = block(x, emb) 480 | x = x.type(self.dtype) 481 | out = self.out_conv(silu(self.out_norm(x))) 482 | return out 483 | 484 | 485 | class SongUNet(torch.nn.Module): 486 | def __init__(self, 487 | image_size, 488 | in_channels, 489 | model_channels, 490 | out_channels, 491 | num_res_blocks, 492 | attention_resolutions=(16), 493 | dropout=0, 494 | channel_mult=(1, 2, 4, 8), 495 | conv_resample=True, 496 | dims=2, 497 | num_classes=None, 498 | use_checkpoint=False, 499 | use_fp16=False, 500 | num_heads=1, 501 | num_head_channels=-1, 502 | num_heads_upsample=-1, 503 | use_scale_shift_norm=False, 504 | resblock_updown=False, 505 | use_new_attention_order=False, 506 | random_init=False, 507 | ): 508 | super().__init__() 509 | self.in_channels = in_channels 510 | self.image_size = image_size 511 | self.num_classes = num_classes 512 | self.random_init = random_init 513 | self.dtype = torch.float16 if use_fp16 else torch.float32 514 | 515 | img_resolution=image_size # Image resolution at input/output. 516 | in_channels=in_channels # Number of color channels at input. 517 | out_channels=out_channels # Number of color channels at output. 518 | label_dim = num_classes if num_classes else 0 # Number of class labels, 0 = unconditional. 519 | augment_dim = 0 # Augmentation label dimensionality, 0 = no augmentation. 520 | 521 | model_channels = model_channels # Base multiplier for the number of channels. 522 | channel_mult = channel_mult # Per-resolution multipliers for the number of channels. 523 | channel_mult_emb = 4 # Multiplier for the dimensionality of the embedding vector. 524 | num_blocks = num_res_blocks # Number of residual blocks per resolution. 525 | attn_resolutions = attention_resolutions # List of resolutions with self-attention. 526 | dropout = 0.0 # Dropout probability of intermediate activations. 527 | label_dropout = 0 # Dropout probability of class labels for classifier-free guidance. 528 | 529 | embedding_type = 'positional' # Timestep embedding type: 'positional' for DDPM++, 'fourier' for NCSN++. 530 | channel_mult_noise = 1 # Timestep embedding size: 1 for DDPM++, 2 for NCSN++. 531 | encoder_type = 'standard' # Encoder architecture: 'standard' for DDPM++, 'residual' for NCSN++. 532 | decoder_type = 'standard' # Decoder architecture: 'standard' for both DDPM++ and NCSN++. 533 | resample_filter = [1,1] # Resampling filter: [1,1] for DDPM++, [1,3,3,1] for NCSN++ 534 | 535 | self.label_dropout = label_dropout 536 | emb_channels = model_channels * channel_mult_emb 537 | noise_channels = model_channels * channel_mult_noise 538 | init = dict(init_mode='xavier_uniform') 539 | init_zero = dict(init_mode='xavier_uniform', init_weight=1e-5) 540 | init_attn = dict(init_mode='xavier_uniform', init_weight=np.sqrt(0.2)) 541 | block_kwargs = dict( 542 | emb_channels=emb_channels, num_heads=1, dropout=dropout, skip_scale=np.sqrt(0.5), eps=1e-6, 543 | resample_filter=resample_filter, resample_proj=True, adaptive_scale=False, 544 | init=init, init_zero=init_zero, init_attn=init_attn, 545 | ) 546 | 547 | # Mapping. 548 | self.map_noise = PositionalEmbedding(num_channels=noise_channels, endpoint=True) if embedding_type == 'positional' else FourierEmbedding(num_channels=noise_channels) 549 | self.map_label = Linear(in_features=label_dim, out_features=noise_channels, **init) if label_dim else None 550 | self.map_augment = Linear(in_features=augment_dim, out_features=noise_channels, bias=False, **init) if augment_dim else None 551 | self.map_layer0 = Linear(in_features=noise_channels, out_features=emb_channels, **init) 552 | self.map_layer1 = Linear(in_features=emb_channels, out_features=emb_channels, **init) 553 | 554 | if self.random_init: 555 | self.map_init_cond_layer0 = Linear(in_features=noise_channels, out_features=emb_channels, **init) 556 | self.map_init_cond_layer1 = Linear(in_features=emb_channels, out_features=emb_channels, **init) 557 | 558 | # Encoder. 559 | self.enc = torch.nn.ModuleDict() 560 | cout = in_channels 561 | caux = in_channels 562 | for level, mult in enumerate(channel_mult): 563 | res = img_resolution >> level 564 | if level == 0: 565 | cin = cout 566 | cout = model_channels 567 | self.enc[f'{res}x{res}_conv'] = Conv2d(in_channels=cin, out_channels=cout, kernel=3, **init) 568 | else: 569 | self.enc[f'{res}x{res}_down'] = UNetBlock(in_channels=cout, out_channels=cout, down=True, **block_kwargs) 570 | if encoder_type == 'skip': 571 | self.enc[f'{res}x{res}_aux_down'] = Conv2d(in_channels=caux, out_channels=caux, kernel=0, down=True, resample_filter=resample_filter) 572 | self.enc[f'{res}x{res}_aux_skip'] = Conv2d(in_channels=caux, out_channels=cout, kernel=1, **init) 573 | if encoder_type == 'residual': 574 | self.enc[f'{res}x{res}_aux_residual'] = Conv2d(in_channels=caux, out_channels=cout, kernel=3, down=True, resample_filter=resample_filter, fused_resample=True, **init) 575 | caux = cout 576 | for idx in range(num_blocks): 577 | cin = cout 578 | cout = model_channels * mult 579 | attn = (res in attn_resolutions) 580 | self.enc[f'{res}x{res}_block{idx}'] = UNetBlock(in_channels=cin, out_channels=cout, attention=attn, **block_kwargs) 581 | skips = [block.out_channels for name, block in self.enc.items() if 'aux' not in name] 582 | 583 | # Decoder. 584 | self.dec = torch.nn.ModuleDict() 585 | for level, mult in reversed(list(enumerate(channel_mult))): 586 | res = img_resolution >> level 587 | if level == len(channel_mult) - 1: 588 | self.dec[f'{res}x{res}_in0'] = UNetBlock(in_channels=cout, out_channels=cout, attention=True, **block_kwargs) 589 | self.dec[f'{res}x{res}_in1'] = UNetBlock(in_channels=cout, out_channels=cout, **block_kwargs) 590 | else: 591 | self.dec[f'{res}x{res}_up'] = UNetBlock(in_channels=cout, out_channels=cout, up=True, **block_kwargs) 592 | for idx in range(num_blocks + 1): 593 | cin = cout + skips.pop() 594 | cout = model_channels * mult 595 | attn = (idx == num_blocks and res in attn_resolutions) 596 | self.dec[f'{res}x{res}_block{idx}'] = UNetBlock(in_channels=cin, out_channels=cout, attention=attn, **block_kwargs) 597 | if decoder_type == 'skip' or level == 0: 598 | if decoder_type == 'skip' and level < len(channel_mult) - 1: 599 | self.dec[f'{res}x{res}_aux_up'] = Conv2d(in_channels=out_channels, out_channels=out_channels, kernel=0, up=True, resample_filter=resample_filter) 600 | self.dec[f'{res}x{res}_aux_norm'] = GroupNorm(num_channels=cout, eps=1e-6) 601 | self.dec[f'{res}x{res}_aux_conv'] = Conv2d(in_channels=cout, out_channels=out_channels, kernel=3, **init_zero) 602 | 603 | def convert_to_fp16(self): 604 | """ 605 | Convert the torso of the model to float16. 606 | """ 607 | self.enc.apply(convert_module_to_f16) 608 | # self.middle_block.apply(convert_module_to_f16) 609 | self.dec.apply(convert_module_to_f16) 610 | 611 | def convert_to_fp32(self): 612 | """ 613 | Convert the torso of the model to float32. 614 | """ 615 | self.enc.apply(convert_module_to_f32) 616 | # self.middle_block.apply(convert_module_to_f32) 617 | self.dec.apply(convert_module_to_f32) 618 | 619 | def forward(self, x, timesteps, y=None): 620 | # Mapping. 621 | class_labels = torch.nn.functional.one_hot(y, num_classes=self.num_classes).to(x.dtype) if y is not None else None 622 | if not isinstance(timesteps, list): 623 | timesteps = [timesteps] 624 | 625 | noise_labels = timesteps[0] 626 | 627 | # Mapping. 628 | emb = self.map_noise(noise_labels) 629 | emb = emb.reshape(emb.shape[0], 2, -1).flip(1).reshape(*emb.shape) # swap sin/cos 630 | if self.map_label is not None: 631 | tmp = class_labels 632 | if self.training and self.label_dropout: 633 | tmp = tmp * (torch.rand([x.shape[0], 1], device=x.device) >= self.label_dropout).to(tmp.dtype) 634 | emb = emb + self.map_label(tmp * np.sqrt(self.map_label.in_features)) 635 | # if self.map_augment is not None and augment_labels is not None: 636 | # emb = emb + self.map_augment(augment_labels) 637 | emb = silu(self.map_layer0(emb)) 638 | emb = silu(self.map_layer1(emb)) 639 | 640 | if self.random_init: 641 | emb_init = self.map_noise(timesteps[1]) 642 | emb_init = emb_init.reshape(emb_init.shape[0], 2, -1).flip(1).reshape(*emb_init.shape) # swap sin/cos 643 | emb_init = silu(self.map_init_cond_layer0(emb_init)) 644 | emb_init = silu(self.map_init_cond_layer1(emb_init)) 645 | emb = emb + emb_init 646 | 647 | # Encoder. 648 | skips = [] 649 | x = x.type(self.dtype) 650 | aux = x 651 | for name, block in self.enc.items(): 652 | if 'aux_down' in name: 653 | aux = block(aux) 654 | elif 'aux_skip' in name: 655 | x = skips[-1] = x + block(aux) 656 | elif 'aux_residual' in name: 657 | x = skips[-1] = aux = (x + block(aux)) / np.sqrt(2) 658 | else: 659 | x = block(x, emb) if isinstance(block, UNetBlock) else block(x) 660 | skips.append(x) 661 | 662 | # Decoder. 663 | aux = None 664 | tmp = None 665 | x = x.type(self.dtype) 666 | for name, block in self.dec.items(): 667 | if 'aux_up' in name: 668 | aux = block(aux) 669 | elif 'aux_norm' in name: 670 | tmp = block(x) 671 | elif 'aux_conv' in name: 672 | tmp = block(silu(tmp)) 673 | aux = tmp if aux is None else tmp + aux 674 | else: 675 | if x.shape[1] != block.in_channels: 676 | x = torch.cat([x, skips.pop()], dim=1) 677 | x = block(x, emb) 678 | return aux 679 | 680 | # @persistence.persistent_class 681 | class EDMPrecond(torch.nn.Module): 682 | def __init__(self, 683 | img_resolution, # Image resolution. 684 | img_channels, # Number of color channels. 685 | label_dim = 0, # Number of class labels, 0 = unconditional. 686 | use_fp16 = False, # Execute the underlying model at FP16 precision? 687 | sigma_min = 0, # Minimum supported noise level. 688 | sigma_max = float('inf'), # Maximum supported noise level. 689 | sigma_data = 0.5, # Expected standard deviation of the training data. 690 | model_type = 'DhariwalUNet', # Class name of the underlying model. 691 | **model_kwargs, # Keyword arguments for the underlying model. 692 | ): 693 | super().__init__() 694 | self.img_resolution = img_resolution 695 | self.img_channels = img_channels 696 | self.label_dim = label_dim 697 | self.use_fp16 = use_fp16 698 | self.sigma_min = sigma_min 699 | self.sigma_max = sigma_max 700 | self.sigma_data = sigma_data 701 | self.model = globals()[model_type](img_resolution=img_resolution, in_channels=img_channels, out_channels=img_channels, label_dim=label_dim, **model_kwargs) 702 | 703 | def forward(self, x, sigma, class_labels=None, force_fp32=False, **model_kwargs): 704 | # x = x.to(torch.float32) 705 | sigma = sigma.reshape(-1, 1, 1, 1) 706 | class_labels = None if self.label_dim == 0 else torch.zeros([1, self.label_dim], device=x.device) if class_labels is None else class_labels.reshape(-1, self.label_dim) 707 | dtype = torch.float16 if (self.use_fp16 and not force_fp32 and x.device.type == 'cuda') else torch.float32 708 | 709 | c_skip = self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2) 710 | c_out = sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2).sqrt() 711 | c_in = 1 / (self.sigma_data ** 2 + sigma ** 2).sqrt() 712 | c_noise = sigma.log() / 4 713 | 714 | F_x = self.model((c_in * x), c_noise.flatten(), class_labels=class_labels, **model_kwargs) 715 | assert F_x.dtype == dtype 716 | D_x = c_skip * x + c_out * F_x 717 | return D_x 718 | 719 | def round_sigma(self, sigma): 720 | return torch.as_tensor(sigma) 721 | 722 | #---------------------------------------------------------------------------- 723 | --------------------------------------------------------------------------------