Showing preview only (234K chars total). Download the full file or copy to clipboard to get everything.
Repository: sangyun884/blur-diffusion
Branch: main
Commit: 9dd8adb917bf
Files: 27
Total size: 224.1 KB
Directory structure:
gitextract_8zm764cm/
├── EMA.py
├── LICENSE
├── README.md
├── blur_diffusion.py
├── eval_x0hat.py
├── eval_x0hat.sh
├── fid.py
├── guided_diffusion/
│ ├── __init__.py
│ ├── dist_util.py
│ ├── fp16_util.py
│ ├── gaussian_diffusion.py
│ ├── image_datasets.py
│ ├── logger.py
│ ├── losses.py
│ ├── nn.py
│ ├── resample.py
│ ├── respace.py
│ ├── script_util.py
│ ├── train_util.py
│ └── unet.py
├── main.py
├── pytorch_fid/
│ ├── __init__.py
│ ├── __main__.py
│ ├── fid_score.py
│ └── inception.py
├── train.sh
└── utils.py
================================================
FILE CONTENTS
================================================
================================================
FILE: EMA.py
================================================
# ---------------------------------------------------------------
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
#
# This work is licensed under the NVIDIA Source Code License
# for Denoising Diffusion GAN. To view a copy of this license, see the LICENSE file.
# ---------------------------------------------------------------
'''
Codes adapted from https://github.com/NVlabs/LSGM/blob/main/util/ema.py
'''
import warnings
import torch
from torch.optim import Optimizer
class EMA(Optimizer):
def __init__(self, opt, ema_decay):
self.ema_decay = ema_decay
self.apply_ema = self.ema_decay > 0.
self.optimizer = opt
self.state = opt.state
self.param_groups = opt.param_groups
def step(self, *args, **kwargs):
retval = self.optimizer.step(*args, **kwargs)
# stop here if we are not applying EMA
if not self.apply_ema:
return retval
ema, params = {}, {}
for group in self.optimizer.param_groups:
for i, p in enumerate(group['params']):
if p.grad is None:
continue
state = self.optimizer.state[p]
# State initialization
if 'ema' not in state:
state['ema'] = p.data.clone()
if p.shape not in params:
params[p.shape] = {'idx': 0, 'data': []}
ema[p.shape] = []
params[p.shape]['data'].append(p.data)
ema[p.shape].append(state['ema'])
for i in params:
params[i]['data'] = torch.stack(params[i]['data'], dim=0)
ema[i] = torch.stack(ema[i], dim=0)
ema[i].mul_(self.ema_decay).add_(params[i]['data'], alpha=1. - self.ema_decay)
for p in group['params']:
if p.grad is None:
continue
idx = params[p.shape]['idx']
self.optimizer.state[p]['ema'] = ema[p.shape][idx, :]
params[p.shape]['idx'] += 1
return retval
def load_state_dict(self, state_dict):
super(EMA, self).load_state_dict(state_dict)
# load_state_dict loads the data to self.state and self.param_groups. We need to pass this data to
# the underlying optimizer too.
self.optimizer.state = self.state
self.optimizer.param_groups = self.param_groups
def swap_parameters_with_ema(self, store_params_in_ema):
""" This function swaps parameters with their ema values. It records original parameters in the ema
parameters, if store_params_in_ema is true."""
# stop here if we are not applying EMA
if not self.apply_ema:
warnings.warn('swap_parameters_with_ema was called when there is no EMA weights.')
return
for group in self.optimizer.param_groups:
for i, p in enumerate(group['params']):
if not p.requires_grad:
continue
ema = self.optimizer.state[p]['ema']
if store_params_in_ema:
tmp = p.data.detach()
p.data = ema.detach()
self.optimizer.state[p]['ema'] = tmp
else:
p.data = ema.detach()
================================================
FILE: LICENSE
================================================
MIT License
Copyright (c) 2022 Sangyun Lee
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
================================================
FILE: README.md
================================================
# blur-diffusion
This is the codebase for our paper Progressive Deblurring of Diffusion Models for Coarse-to-Fine Image Synthesis.
> **Progressive Deblurring of Diffusion Models for Coarse-to-Fine Image Synthesis**<br>
> Sangyun Lee<sup>1</sup>, Hyungjin Chung<sup>2</sup>, Jaehyeon Kim<sup>3</sup>, Jong Chul Ye<sup>2</sup>
> <sup>1</sup>Soongsil University, <sup>2</sup>KAIST, <sup>3</sup>Kakao Enterprise<br>
> Paper: https://arxiv.org/abs/2207.11192<br>
> **Abstract:** *Recently, diffusion models have shown remarkable results in image synthesis by gradually removing noise and amplifying signals.
Although the simple generative process surprisingly works well, is this the best way to generate image data? For instance, despite the fact that human perception is more sensitive to the low-frequencies of an image, diffusion models themselves do not consider any relative importance of each frequency component. Therefore, to incorporate the inductive bias for image data, we propose a novel generative process that synthesizes images in a coarse-to-fine manner. First, we generalize the standard diffusion models by enabling diffusion in a rotated coordinate system with different velocities for each component of the vector. We further propose a blur diffusion as a special case, where each frequency component of an image is diffused at different speeds. Specifically, the proposed blur diffusion consists of a forward process that blurs an image and adds noise gradually, after which a corresponding reverse process deblurs an image and removes noise progressively. Experiments show that proposed model outperforms the previous method in FID on LSUN bedroom and church datasets.*
## Train
```bash train.sh```
## Visualization
```bash eval_x0hat.sh```
## Dataset
Image files are required to compute FID during training.
================================================
FILE: blur_diffusion.py
================================================
import torch
import torch.nn.functional as F
import time
import utils
import os
import numpy as np
from utils import clear
def gaussian_kernel_1d(kernel_size, sigma):
assert sigma > 0.00001
x_cord = torch.arange(kernel_size)
mean = (kernel_size - 1) / 2.
# pdf of 1d gaussian not considering normalization constant
gaussian_kernel = torch.exp(-0.5 * ((x_cord - mean)/sigma)**2)
gaussian_kernel = gaussian_kernel / torch.sum(gaussian_kernel)
return gaussian_kernel
def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
"""
Ported from https://github.com/openai/guided-diffusion
Utility function for cosine noise schedule
"""
betas = []
for i in range(num_diffusion_timesteps):
t1 = i / num_diffusion_timesteps
t2 = (i + 1) / num_diffusion_timesteps
betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
return np.array(betas)
class ExpSchedule:
def __init__(self, N, offset = 1e-4):
self.N = N
def f(i):
#return 1 - np.sin((i / N + 1) * np.pi / 2) + offset
return np.exp(5*i/N - 5) + offset
idxs = np.arange(N+1)
self.alphas_bar = 1 - f(idxs) / f(idxs[-1])
self.alphas_bar_left_shifted = 1 - f(idxs-1) / f(idxs[-1])
self.alphas = self.alphas_bar / self.alphas_bar_left_shifted
self.betas = 1 - self.alphas
def get_betas(self):
return self.betas
# Increase blur strength as i increases
class ForwardBlurIncreasing:
def __init__(self, N, beta_min, beta_max, sig, sig_min, sig_max, D_diag, blur=None, noise_schedule='linear',
channel=3, resolution=32, pad=None, device='cuda:0', f_type = 'linear'):
# N: total number of discretizations used
self.device = device
self.N = N
self.beta_min = beta_min
self.beta_max = beta_max
self.sig = sig
self.sig_min = sig_min # \sigma_1
self.sig_max = sig_max # \sigma_N
self.D_diag = D_diag
# total dimension of image
self.dim = self.D_diag.shape[0]
# "deblur" class
self.blur = blur
self.resolution = resolution
self.channel = channel
self.pad = pad
self.noise_schedule = noise_schedule
if noise_schedule == 'linear':
self.betas = torch.linspace(beta_min, beta_max, N, device=device)
elif noise_schedule == 'cosine':
self.betas = torch.tensor(betas_for_alpha_bar(
N,
lambda t: np.cos((t + 0.008) / 1.008 * np.pi / 2) ** 2,
), device=device, dtype=torch.float)
elif noise_schedule == 'exp':
Exp = ExpSchedule(N)
betas = torch.tensor(Exp.get_betas())
self.betas = torch.tensor(betas, device=device)
self.betas = F.pad(self.betas, (1, 0), value=0.0)
self.alphas = 1 - self.betas
self.sqrt_alphas = torch.sqrt(self.alphas)
self.alphas_bar = torch.cumprod(self.alphas, dim=0)
self.sqrt_alphas_bar = torch.cumprod(self.sqrt_alphas, dim=0)
self.alphas_bar_prev = F.pad(self.alphas_bar, [1, 0], value=1)[:N]
# def f(i):
# a = (self.sig_max**2 - self.sig_min**2) / (self.sig**2 * (self.N**2 - 1))
# b = (self.sig_min / self.sig) ** 2 - a
# return a * i**2 + b
def f_linear(i):
f_N = (self.sig_max / self.sig) ** 2
f_1 = (self.sig_min / self.sig) ** 2
return (f_N - f_1)/ (N-1) * (i - 1) + f_1
idxs = torch.tensor(np.arange(N+1), device=device, dtype=torch.float)
def f_log(i):
log = lambda x: torch.log(x + 1e-6) / (10*np.log(N))
f_N = (self.sig_max / self.sig) ** 2
f_1 = (self.sig_min / self.sig) ** 2
a = (f_N - f_1) / log(torch.tensor(N))
b = f_1
return a*log(i) + b
def f_quadratic(i):
f_N = (sig_max / sig) ** 2
f_1 = (sig_min / sig) ** 2
a = (f_N - f_1) / (N**2 - 1)
b = f_1 - (f_N - f_1) / (N**2 - 1)
return a*i**2 + b
def cubic(i):
f_N = (sig_max / sig) ** 2
f_1 = (sig_min / sig) ** 2
a = (f_N-f_1) / N**3
b = f_1
return a*i**3 + b
def quartic(i):
f_N = (sig_max / sig) ** 2
f_1 = (sig_min / sig) ** 2
a = (f_N-f_1) / N**4
b = f_1
return a*i**4 + b
def triangular(i):
less_than_N_2 = (i < N/2).type(torch.int)
return f_linear(i) * less_than_N_2 + f_linear(N-i) * (1-less_than_N_2 )
if f_type == 'linear':
f = f_linear
elif f_type == 'log':
f = f_log
elif f_type == 'quadratic':
f = f_quadratic
elif f_type == 'cubic':
f = cubic
elif f_type == 'quartic':
f = quartic
elif f_type == 'triangular':
f = triangular
else:
raise NotImplementedError
self.fs = f(idxs)
print("fs: ", self.fs)
self.fs_cum = torch.cumsum(self.fs, dim=0)
idxs_long = torch.tensor(np.arange(N+1), device=device, dtype=torch.long)
def B(i):
p = (2 * f(i)).unsqueeze(-1).repeat(1, self.D_diag.shape[0])
print("p: ", p.shape)
D = self.D_diag.repeat(p.shape[0], 1)
print("D: ", D.shape)
return self.alphas[i].unsqueeze(-1) * self.D_diag ** p
self.Bs = B(idxs_long)
self.Bs_bar = F.pad(torch.cumprod(self.Bs[1:], dim=0), (0,0,1, 0), value=0)
self.one_minus_Bs_bar = 1 - self.Bs_bar
self.one_minus_Bs_bar_sqrt = torch.sqrt(self.one_minus_Bs_bar)
self.Bs_sqrt = torch.sqrt(self.Bs)
self.Bs_squared = self.Bs ** 2
self.Bs_bar_sqrt = torch.sqrt(self.Bs_bar)
# fname = f'res{resolution}_N{N}_sig{sig}_sigmin_{sig_min}_sigmax_{sig_max}_b_mat_sq-nc{channel}_{noise_schedule}_inc_nopad_rotated.npz'
# if os.path.isfile(fname):
# with np.load(fname, allow_pickle=True) as data:
# print("data", list(data.keys()))
# data = data['arr'].astype(np.float32)
# b_mat_sq = torch.tensor(data, device=device)
# self.b_mat_sq = b_mat_sq
# else:
# tic = time.time()
# self.b_mat_sq = torch.zeros([N + 1, self.dim]).to(device)
# self.b_mat = torch.sqrt(self.b_mat_sq)
def _beta_i(self, i):
return self.betas[i]
def _alpha_i(self, i):
return self.alphas[i]
def _alphas_bar_i(self, i):
return self.alphas_bar[i]
def _Bhat_i(self, i):
return self.b_mat[i, :]
# Forward blur
def _Bhat_sq_i(self, i):
assert i != 0
return self.b_mat_sq[i, :]
# Forward blur
def get_mean(self, x0, i):
mat = self.Bs_bar_sqrt[i]
mean = self.blur.U(mat * self.blur.Ut(x0))
return mean
# Forward blur
def get_std(self,i, noise):
mat = self.one_minus_Bs_bar_sqrt[i]
std = self.blur.U(mat * self.blur.Ut(noise))
return std
# Forward blur
def get_var(self, x0, i, noise=None):
raise NotImplementedError
Bhat_i = self._Bhat_sq_i(i)
noise = noise if noise is not None else torch.randn_like(x0)
std = self.blur.U(Bhat_i * self.blur.Ut(noise))
return std
def W(self, x, i):
mat = self.Bs_sqrt[i]
batch = x.shape[0]
blurred = self.blur.U(mat * self.blur.Ut(x)).view(batch,
self.channel,
self.resolution,
self.resolution)
return blurred
def W_inv(self, x, i):
mat = self.Bs_squared[i]
batch = x.shape[0]
blurred = self.blur.U(mat * self.blur.Ut(x)).view(batch,
self.channel,
self.resolution,
self.resolution)
return blurred
def U_I_minus_B_Ut(self, x, i): # U(I-B)UT
batch = x.shape[0]
mat = 1 - self.Bs[i]
result = self.blur.U(mat * self.blur.Ut(x)).view(batch,
self.channel,
self.resolution,
self.resolution)
return result
def U_I_minus_B_sqrt_Ut(self, x, i): # U(sqrt(I-B))UT
batch = x.shape[0]
mat = torch.sqrt(1 - self.Bs[i])
result = self.blur.U(mat * self.blur.Ut(x)).view(batch,
self.channel,
self.resolution,
self.resolution)
return result
# Forward blur
def get_x_i(self, x0, i, return_eps = False):
assert 0 not in i
if len(x0.shape) == 4:
batch = x0.shape[0]
# reparam trick
mean = self.get_mean(x0, i)
noise = torch.randn_like(x0)
std = self.get_std(i, noise = noise)
# crop
if len(x0.shape) == 4:
img = (mean + std).view(batch,
self.channel,
self.resolution,
self.resolution)
else:
img = (mean + std).view(self.channel,
self.resolution,
self.resolution)
# cropped = cropped + torch.randn_like(cropped) * torch.sqrt(self.betas[i-1])
if return_eps:
return img, noise
return img
def get_x_N(self, x0_shape, N):
"""
prior sampling
"""
if len(x0_shape) == 4:
batch = x0_shape[0]
# pad
x0 = torch.zeros(x0_shape, device=self.device)
noise = torch.randn_like(x0)
std = self.get_std(N, noise)
# crop
if len(x0_shape) == 4:
img = std.view(batch,
self.channel,
self.resolution,
self.resolution)
else:
img = std.view(self.channel,
self.resolution,
self.resolution)
return img
def get_x0_from_eps(self, xi, eps, i):
batch = xi.shape[0]
std = self.get_std(i, noise = eps).view(batch,
self.channel,
self.resolution,
self.resolution)
mean = xi - std
return mean / torch.sqrt(self.alphas_bar[i]).view(-1, 1, 1, 1)
mat = self.Bs_bar_sqrt[i] ** (-1)
x0 = self.blur.U(mat * self.blur.Ut(mean)).view(batch,
self.channel,
self.resolution,
self.resolution)
return x0
def get_score_gt(self, xi, x0, i):
# Return the ground-truth score of x_i
assert torch.all(i>0)
batch = x0.shape[0]
mean = self.get_mean(x0, i).view(batch,
self.channel,
self.resolution,
self.resolution)
diff = xi - mean
mat = (self.one_minus_Bs_bar[i])**(-1)
score = -self.blur.U(mat * self.blur.Ut(diff)).view(batch,
self.channel,
self.resolution,
self.resolution)
return score
def get_score_gt2(self, xi, x0, i):
# Return the ground-truth score of x_i
batch = x0.shape[0]
x0_pad = utils.pad(x0, self.pad)
# a_iW^ix_0
mean_pad = self.get_mean(x0_pad, i).view(batch,
self.channel,
self.resolution + self.pad * 2,
self.resolution + self.pad * 2)
mean = utils.crop(mean_pad, self.pad)
diff = xi - mean
diff = utils.pad(diff, self.pad)
score = self.blur.U(self._Bhat_i(i) ** (-2) * self.blur.Ut(diff)).view(batch,
self.channel,
self.resolution + self.pad * 2,
self.resolution + self.pad * 2)
score = -utils.crop(score, self.pad)
return score
def sanity(self, x_0, i):
batch = x_0.shape[0]
tol = 1e-2
xi, eps = self.get_x_i(x_0, i, return_eps = True)
score1 = self.get_score_from_eps(eps, i)
score2 = self.get_score_gt(xi, x_0, i)
# utils.tensor_imsave(score1[0], "./", "score1.jpg")
# utils.tensor_imsave(score2[0], "./", "score2.jpg")
# rms = lambda x: torch.sqrt(torch.mean(x**2))
# print(f"rms(score1) = {rms(score1)}, rms(score2) = {rms(score2)}")
# noise = torch.randn_like(x_0)
# noise_pad = utils.pad(noise, self.pad)
# std = self.blur.U(self._Bhat_i(i) * self.blur.Ut(noise_pad)).view(batch,
# self.channel,
# self.resolution + self.pad * 2,
# self.resolution + self.pad * 2)
# std = utils.crop(std, self.pad)
# std = utils.pad(std, self.pad)
# # crop -> pad makes error
# score1 = self.blur.U(self._Bhat_i(i) ** (-2) * self.blur.Ut(std)).view(batch,
# self.channel,
# self.resolution + self.pad * 2,
# self.resolution + self.pad * 2)
# score1 = -utils.crop(score1, self.pad)
# score2 = self.blur.U(self._Bhat_i(i) ** (-1) * self.blur.Ut(noise_pad)).view(batch, self.channel, self.resolution + self.pad * 2, self.resolution + self.pad * 2)
# score2 = -utils.crop(score2, self.pad)
MAE = torch.mean(torch.abs(score1 - score2))
assert MAE < tol, f'MAE = {MAE}'
print(f"MAE = {MAE}")
# score3 = self.get_score_gt2(xi, x_0, i)
# print(f"score3 = {score3[:5]}")
# print(f"score2 = {score2[:5]}")
# MAE = torch.mean(torch.abs(score3 - score2))
# print(f"MAE = {MAE}, fraction = {MAE/torch.mean(torch.abs(score3))}, norm: {torch.mean(torch.abs(score3))}")
# assert MAE <1e-5, f'MAE = {MAE}'
# MAE = torch.mean(torch.abs(self.b_mat_sq[1:] - self.one_minus_Bs_bar[1:]))
# assert MAE <1e-5, f'MAE = {MAE}'
def get_score_from_eps(self, eps, i):
# Return the score of x_i
batch = eps.shape[0]
mat = self.one_minus_Bs_bar_sqrt[i]**(-1)
score = -self.blur.U(mat * self.blur.Ut(eps)).view(batch, self.channel, self.resolution, self.resolution)
return score
def get_score_from_std(self, std, i):
# Return the score of x_i
batch = std.shape[0]
mat = self.one_minus_Bs_bar[i] ** (-1)
score = -self.blur.U(mat * self.blur.Ut(std)).view(batch, self.channel, self.resolution, self.resolution)
return score
def get_loss_i_exact(self, model, x0, xi, i):
# Calculate the DSM objective (exact)
batch = x0.shape[0]
pred = model(xi, i)
score = self.get_score_gt(x0, xi, i)
loss = torch.mean((pred - score) ** 2)
return loss
def get_loss_i_simple(self, model, x0, xi, i):
raise NotImplementedError
# Calculate the DSM objective (simple)
batch = x0.shape[0]
pred = model(xi, i)
# UB**2U^Ts
pred = utils.pad(pred, self.pad)
pred = self.blur.U(self._Bhat_i(i) ** 2 * self.blur.Ut(pred)).view(batch,
self.channel,
self.resolution + self.pad * 2,
self.resolution + self.pad * 2)
pred = utils.crop(pred, self.pad)
# ai[UDUT]x0
x0_pad = utils.pad(x0, self.pad)
mean = self.get_mean(x0_pad, i).view(batch,
self.channel,
self.resolution + self.pad * 2,
self.resolution + self.pad * 2)
mean = utils.crop(mean, self.pad)
loss = torch.mean((pred + xi - mean) ** 2)
return loss
def get_loss_i_eps_simple(self, model, x_i, i, eps):
pred = model(x_i, i)
loss = torch.mean((pred - eps) ** 2)
return loss
def get_loss_i_std_matching(self, model, x_i, i, eps):
pred = model(x_i, i)
batch = x_i.shape[0]
std = self.get_std(i, noise = eps).view(batch,
self.channel,
self.resolution,
self.resolution)
loss = torch.mean((pred - std) ** 2)
return loss
class H_functions:
"""
Ported from https://github.com/bahjat-kawar/ddrm
A class replacing the SVD of a matrix H, perhaps efficiently.
All input vectors are of shape (Batch, ...).
All output vectors are of shape (Batch, DataDimension).
"""
def V(self, vec):
"""
Multiplies the input vector by V
"""
raise NotImplementedError()
def Vt(self, vec):
"""
Multiplies the input vector by V transposed
"""
raise NotImplementedError()
def U(self, vec):
"""
Multiplies the input vector by U
"""
raise NotImplementedError()
def Ut(self, vec):
"""
Multiplies the input vector by U transposed
"""
raise NotImplementedError()
def singulars(self):
"""
Returns a vector containing the singular values. The shape of the vector should be the same as the smaller dimension (like U)
"""
raise NotImplementedError()
def add_zeros(self, vec):
"""
Adds trailing zeros to turn a vector from the small dimension (U) to the big dimension (V)
"""
raise NotImplementedError()
def H(self, vec):
"""
Multiplies the input vector by H
"""
temp = self.Vt(vec)
singulars = self.singulars()
return self.U(singulars * temp[:, :singulars.shape[0]])
def Ht(self, vec):
"""
Multiplies the input vector by H transposed
"""
temp = self.Ut(vec)
singulars = self.singulars()
return self.V(self.add_zeros(singulars * temp[:, :singulars.shape[0]]))
def H_pinv(self, vec):
"""
Multiplies the input vector by the pseudo inverse of H
"""
temp = self.Ut(vec)
singulars = self.singulars()
temp[:, :singulars.shape[0]] = temp[:, :singulars.shape[0]] / singulars
return self.V(self.add_zeros(temp))
class Deblurring(H_functions):
def mat_by_img(self, M, v):
return torch.matmul(M, v.reshape(v.shape[0] * self.channels, self.img_dim,
self.img_dim)).reshape(v.shape[0], self.channels, M.shape[0], self.img_dim)
def img_by_mat(self, v, M):
return torch.matmul(v.reshape(v.shape[0] * self.channels, self.img_dim,
self.img_dim), M).reshape(v.shape[0], self.channels, self.img_dim, M.shape[1])
def __init__(self, kernel, channels, img_dim, device):
self.img_dim = img_dim
self.channels = channels
# build 1D conv matrix
H_small = torch.zeros(img_dim, img_dim, device=device)
for i in range(img_dim):
for j in range(i - kernel.shape[0] // 2, i + kernel.shape[0] // 2):
if j < 0 or j >= img_dim: continue
H_small[i, j] = kernel[j - i + kernel.shape[0] // 2]
# torch.cholesky(H_small) # raise error if H_small is not positive definite
self.H_small = H_small # positive definite matrix
# get the evd of the 1D conv
self.U_small, self.singulars_small, _ = torch.svd(H_small, some=False)
self.V_small = self.U_small # V should equal U since H is symmetric
ZERO = 3e-2
# truncate
self.singulars_small[self.singulars_small < ZERO] = ZERO
# calculate the singular values of the big matrix using Kronecker product
self._singulars = torch.matmul(self.singulars_small.reshape(img_dim, 1),
self.singulars_small.reshape(1, img_dim)).reshape(img_dim ** 2)
self._singulars[self._singulars > 1] = 1
print("contains zero?", torch.any(self._singulars == 0))
# raise NotImplementedError
# sort the big matrix singulars and save the permutation
self._singulars, self._perm = self._singulars.sort(descending=True) # , stable=True)
def V(self, vec):
# invert the permutation
temp = torch.zeros(vec.shape[0], self.img_dim ** 2, self.channels, device=vec.device)
temp[:, self._perm, :] = vec.clone().reshape(vec.shape[0], self.img_dim ** 2, self.channels)
temp = temp.permute(0, 2, 1)
# multiply the image by V from the left and by V^T from the right
out = self.mat_by_img(self.V_small, temp)
out = self.img_by_mat(out, self.V_small.transpose(0, 1)).reshape(vec.shape[0], -1)
return out
def Vt(self, vec):
# multiply the image by V^T from the left and by V from the right
temp = self.mat_by_img(self.V_small.transpose(0, 1), vec.clone())
temp = self.img_by_mat(temp, self.V_small).reshape(vec.shape[0], self.channels, -1)
# permute the entries according to the singular values
temp = temp[:, :, self._perm].permute(0, 2, 1)
return temp.reshape(vec.shape[0], -1)
def U(self, vec):
# invert the permutation
temp = torch.zeros(vec.shape[0], self.img_dim ** 2, self.channels, device=vec.device)
temp[:, self._perm, :] = vec.clone().reshape(vec.shape[0], self.img_dim ** 2, self.channels)
temp = temp.permute(0, 2, 1)
# multiply the image by U from the left and by U^T from the right
out = self.mat_by_img(self.U_small, temp)
out = self.img_by_mat(out, self.U_small.transpose(0, 1)).reshape(vec.shape[0], -1)
return out
def Ut(self, vec):
# multiply the image by U^T from the left and by U from the right
temp = self.mat_by_img(self.U_small.transpose(0, 1), vec.clone())
temp = self.img_by_mat(temp, self.U_small).reshape(vec.shape[0], self.channels, -1)
# permute the entries according to the singular values
temp = temp[:, :, self._perm].permute(0, 2, 1)
return temp.reshape(vec.shape[0], -1)
def conv1d_col_matmul(self, x):
return torch.matmul(self.H_small, x)
def conv1d_row_matmul(self, x):
return torch.matmul(x, self.H_small)
def conv2d_sep_matmul(self, x):
return torch.matmul(torch.matmul(self.H_small, x), self.H_small)
def singulars(self):
if self.channels == 3:
return self._singulars.repeat(1, 3).reshape(-1)
else:
return self._singulars.reshape(-1)
def update_singulars(self, new_singulars):
self._singulars = new_singulars
def add_zeros(self, vec):
return vec.clone().reshape(vec.shape[0], -1)
================================================
FILE: eval_x0hat.py
================================================
import torch
from PIL import Image
from collections import defaultdict
import torch.nn.functional as TF
import torchvision.datasets as dsets
from torchvision import transforms
import numpy as np
import torch.optim as optim
import torchvision
import matplotlib.pyplot as plt
import utils
from guided_diffusion.unet import UNetModel
import math
from tensorboardX import SummaryWriter
import os
import json
from collections import namedtuple
import argparse
from torchvision.utils import save_image
from tqdm import tqdm
from blur_diffusion import Deblurring, ForwardBlurIncreasing, gaussian_kernel_1d
from utils import normalize_np, clear
from EMA import EMA
from torch.nn import DataParallel
parser = argparse.ArgumentParser(description='Configs')
parser.add_argument('--gpu', type=str, help='gpu num')
parser.add_argument('--name', type=str, help='Saving directory name')
parser.add_argument('--ckpt', default='', type=str, help='UNet checkpoint')
parser.add_argument('--bsize', default=16, type=int, help='batchsize')
parser.add_argument('--N', default=1000, type=int, help='Max diffusion timesteps')
parser.add_argument('--sig', default=0.3, type=float, help='sigma value for blur kernel')
parser.add_argument('--sig_min', default=0.3, type=float, help='sigma value for blur kernel')
parser.add_argument('--sig_max', default=1.5, type=float, help='sigma value for blur kernel')
parser.add_argument('--noise_schedule', default='linear', type=str, help='Type of noise schedule to use')
parser.add_argument('--betamin', default=0.0001, type=float, help='beta (min). get_score(1) can diverge if this is too low.')
parser.add_argument('--betamax', default=0.02, type=float, help='beta (max)')
parser.add_argument('--res', type=int, help='resolution')
parser.add_argument('--nc', type=int, help='channels')
parser.add_argument('--loss_type', type=str, default = 'sm_simple', choices=['sm_simple', 'eps_simple', 'sm_exact'])
parser.add_argument('--f_type', type=str, default = 'linear', choices=['linear', 'log', 'quadratic', 'quartic'])
parser.add_argument('--dropout', default=0, type=float, help='dropout')
parser.add_argument('--freq_feat', action='store_true', help = "concat Utx_i")
opt = parser.parse_args()
ksize = opt.res * 2 - 1
pad = 0
bsize = opt.bsize
beta_min = opt.betamin
beta_max = opt.betamax
device = torch.device(f'cuda:{opt.gpu}')
# define forward blur
kernel = gaussian_kernel_1d(ksize, opt.sig)
blur = Deblurring(kernel, opt.nc, opt.res, device=device)
print("blur.U_small.shape:", blur.U_small.shape)
D_diag = blur.singulars()
fb = ForwardBlurIncreasing(N=opt.N, beta_min=beta_min, beta_max=beta_max, sig=opt.sig, sig_max = opt.sig_max, sig_min = opt.sig_min, D_diag=D_diag,
blur=blur, channel=opt.nc, device=device, noise_schedule=opt.noise_schedule, resolution=opt.res, pad=pad, f_type=opt.f_type)
iter = opt.ckpt.split('/')[-1].split('.')[0]
dir = os.path.join('experiments', opt.name, f'inferences-{iter}')
if not os.path.exists(dir):
os.mkdir(dir)
model = UNetModel(opt.res, opt.nc, 128, opt.nc, blur = blur, dropout=opt.dropout, freq_feat = opt.freq_feat)
model.load_state_dict(torch.load(opt.ckpt))
model.to(device)
model.eval()
print("input_nc", opt.nc, "resolution", opt.res)
num_samples = bsize
id = 0
for _ in tqdm(range(num_samples // bsize)):
with torch.no_grad():
i = np.array([opt.N - 1] * bsize)
i = torch.from_numpy(i).to(device)
pred = fb.get_x_N([bsize, opt.nc, opt.res, opt.res], i)
x0_list = []
for i in reversed(range(1, opt.N)):
i = np.array([i] * bsize)
i = torch.from_numpy(i).to(device)
if opt.loss_type == "sm_simple":
s = model(pred, i)
elif opt.loss_type == "eps_simple":
eps = model(pred, i)
s = fb.get_score_from_eps(eps, i)
x0_hat = fb.get_x0_from_eps(pred, eps, i)
if i[0] % 100 == 0 or i[0] == 1:
x0_list.append(x0_hat)
elif opt.loss_type == "sm_exact":
s = model(pred, i)
else:
raise NotImplementedError
s = fb.U_I_minus_B_Ut(s, i)
rms = lambda x: torch.sqrt(torch.mean(x ** 2))
print(f"rms(s) * fb._beta_i(i) = {rms(s) * fb._beta_i(i)[0]}")
hf = pred - fb.W(pred, i)
# Anderson theorem
pred1 = pred + hf # unsharpening mask filtering
pred2 = pred1 + s # # denoising
if i[0] > 2:
pred = pred2 + fb.U_I_minus_B_sqrt_Ut(torch.randn_like(pred), i) # inject noise
else:
pred = pred2
print(f"i = {i[0]}, rmse = {torch.sqrt(torch.mean(pred**2))}, mean = {torch.mean(pred)} std = {torch.std(pred)}" )
grid = torch.cat(x0_list, dim=3) # grid_sample.shape: (bsize, channel, H, W * 12)
# (batch_size, channel, H, W * 12) -> (channel, H * bsize, W * 12)
grid = grid.permute(1, 0, 2, 3).contiguous().view(grid.shape[1], -1, grid.shape[3])
# (bsize, channel, H, W) -> (channel, H, W * bsize)
save_image(grid, os.path.join(dir, f'xhat.png'))
================================================
FILE: eval_x0hat.sh
================================================
#!/bin/bash
export OMP_NUM_THREADS=1
export CUDA_VISIBLE_DEVICES=1
python eval_x0hat.py \
--name test \
--noise_schedule linear \
--gpu 0 \
--bsize 4 \
--sig 0.4 \
--sig_min 0 \
--sig_max 0.15 \
--f_type quartic \
--res 64 \
--nc 3 \
--loss_type eps_simple \
--ckpt /home/ubuntu/code/sangyoon/forward-blur-new/experiments/lsun-simplified-64-sigmax0.15-quartic/model_450001.ckpt
================================================
FILE: fid.py
================================================
from pytorch_fid.fid_score import compute_statistics_of_path, calculate_frechet_distance
from pytorch_fid.inception import InceptionV3
import os
class FID(object):
def __init__(self, real_dir, device, bsize = 128):
# https://github.com/mseitzer/pytorch-fid
# If real_dir contains a .npz file, then we'll just use that.
# Otherwise it's assumed to be a directory.
# In that case, statistics will be computed, and .npz file will be saved.
if os.path.exists(os.path.join(real_dir, "inception_statistics.npz")):
real_dir = os.path.join(real_dir, "inception_statistics.npz")
self.dims = 2048
self.bsize = bsize
self.device = device
block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[self.dims]
self.model = InceptionV3([block_idx]).to(device)
m, s = compute_statistics_of_path(real_dir, self.model, bsize, self.dims, device)
self.mu_real = m
self.sig_real = s
def __call__(self, fake_dir):
m, s = compute_statistics_of_path(fake_dir, self.model, self.bsize, self.dims, self.device)
fid_value = calculate_frechet_distance(m, s, self.mu_real, self.sig_real)
return fid_value
================================================
FILE: guided_diffusion/__init__.py
================================================
"""
Codebase for "Improved Denoising Diffusion Probabilistic Models".
"""
================================================
FILE: guided_diffusion/dist_util.py
================================================
"""
Helpers for distributed training.
"""
import io
import os
import socket
import blobfile as bf
from mpi4py import MPI
import torch as th
import torch.distributed as dist
# Change this to reflect your cluster layout.
# The GPU for a given rank is (rank % GPUS_PER_NODE).
GPUS_PER_NODE = 8
SETUP_RETRY_COUNT = 3
def setup_dist():
"""
Setup a distributed process group.
"""
if dist.is_initialized():
return
os.environ["CUDA_VISIBLE_DEVICES"] = f"{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}"
comm = MPI.COMM_WORLD
backend = "gloo" if not th.cuda.is_available() else "nccl" # backend = "nccl"
if backend == "gloo":
hostname = "localhost"
else:
hostname = socket.gethostbyname(socket.getfqdn())
# print(hostname)
os.environ["MASTER_ADDR"] = comm.bcast(hostname, root=0)
os.environ["RANK"] = str(comm.rank)
os.environ["WORLD_SIZE"] = str(comm.size)
port = comm.bcast(_find_free_port(), root=0)
os.environ["MASTER_PORT"] = str(port)
# print(port)
# print(os.environ["MASTER_ADDR"], os.environ["RANK"], os.environ["WORLD_SIZE"])
print(r"dist.init_process_group(backend=backend, init_method=\"env://\") -- start")
dist.init_process_group(backend=backend, init_method="env://")
print(r"dist.init_process_group(backend=backend, init_method=\"env://\") -- end")
def dev():
"""
Get the device to use for torch.distributed.
"""
if th.cuda.is_available():
return th.device(f"cuda")
return th.device("cpu")
def load_state_dict(path, **kwargs):
"""
Load a PyTorch file without redundant fetches across MPI ranks.
"""
chunk_size = 2 ** 30 # MPI has a relatively small size limit
if MPI.COMM_WORLD.Get_rank() == 0:
with bf.BlobFile(path, "rb") as f:
data = f.read()
num_chunks = len(data) // chunk_size
if len(data) % chunk_size:
num_chunks += 1
MPI.COMM_WORLD.bcast(num_chunks)
for i in range(0, len(data), chunk_size):
MPI.COMM_WORLD.bcast(data[i : i + chunk_size])
else:
num_chunks = MPI.COMM_WORLD.bcast(None)
data = bytes()
for _ in range(num_chunks):
data += MPI.COMM_WORLD.bcast(None)
return th.load(io.BytesIO(data), **kwargs)
def sync_params(params):
"""
Synchronize a sequence of Tensors across ranks from rank 0.
"""
for p in params:
with th.no_grad():
dist.broadcast(p, 0)
def _find_free_port():
try:
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.bind(("", 0))
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
return s.getsockname()[1]
finally:
s.close()
================================================
FILE: guided_diffusion/fp16_util.py
================================================
"""
Helpers to train with 16-bit precision.
"""
import numpy as np
import torch as th
import torch.nn as nn
from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
from . import logger
INITIAL_LOG_LOSS_SCALE = 20.0
def convert_module_to_f16(l):
"""
Convert primitive modules to float16.
"""
if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
l.weight.data = l.weight.data.half()
if l.bias is not None:
l.bias.data = l.bias.data.half()
def convert_module_to_f32(l):
"""
Convert primitive modules to float32, undoing convert_module_to_f16().
"""
if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
l.weight.data = l.weight.data.float()
if l.bias is not None:
l.bias.data = l.bias.data.float()
def make_master_params(param_groups_and_shapes):
"""
Copy model parameters into a (differently-shaped) list of full-precision
parameters.
"""
master_params = []
for param_group, shape in param_groups_and_shapes:
master_param = nn.Parameter(
_flatten_dense_tensors(
[param.detach().float() for (_, param) in param_group]
).view(shape)
)
master_param.requires_grad = True
master_params.append(master_param)
return master_params
def model_grads_to_master_grads(param_groups_and_shapes, master_params):
"""
Copy the gradients from the model parameters into the master parameters
from make_master_params().
"""
for master_param, (param_group, shape) in zip(
master_params, param_groups_and_shapes
):
master_param.grad = _flatten_dense_tensors(
[param_grad_or_zeros(param) for (_, param) in param_group]
).view(shape)
def master_params_to_model_params(param_groups_and_shapes, master_params):
"""
Copy the master parameter data back into the model parameters.
"""
# Without copying to a list, if a generator is passed, this will
# silently not copy any parameters.
for master_param, (param_group, _) in zip(master_params, param_groups_and_shapes):
for (_, param), unflat_master_param in zip(
param_group, unflatten_master_params(param_group, master_param.view(-1))
):
param.detach().copy_(unflat_master_param)
def unflatten_master_params(param_group, master_param):
return _unflatten_dense_tensors(master_param, [param for (_, param) in param_group])
def get_param_groups_and_shapes(named_model_params):
named_model_params = list(named_model_params)
scalar_vector_named_params = (
[(n, p) for (n, p) in named_model_params if p.ndim <= 1],
(-1),
)
matrix_named_params = (
[(n, p) for (n, p) in named_model_params if p.ndim > 1],
(1, -1),
)
return [scalar_vector_named_params, matrix_named_params]
def master_params_to_state_dict(
model, param_groups_and_shapes, master_params, use_fp16
):
if use_fp16:
state_dict = model.state_dict()
for master_param, (param_group, _) in zip(
master_params, param_groups_and_shapes
):
for (name, _), unflat_master_param in zip(
param_group, unflatten_master_params(param_group, master_param.view(-1))
):
assert name in state_dict
state_dict[name] = unflat_master_param
else:
state_dict = model.state_dict()
for i, (name, _value) in enumerate(model.named_parameters()):
assert name in state_dict
state_dict[name] = master_params[i]
return state_dict
def state_dict_to_master_params(model, state_dict, use_fp16):
if use_fp16:
named_model_params = [
(name, state_dict[name]) for name, _ in model.named_parameters()
]
param_groups_and_shapes = get_param_groups_and_shapes(named_model_params)
master_params = make_master_params(param_groups_and_shapes)
else:
master_params = [state_dict[name] for name, _ in model.named_parameters()]
return master_params
def zero_master_grads(master_params):
for param in master_params:
param.grad = None
def zero_grad(model_params):
for param in model_params:
# Taken from https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.add_param_group
if param.grad is not None:
param.grad.detach_()
param.grad.zero_()
def param_grad_or_zeros(param):
if param.grad is not None:
return param.grad.data.detach()
else:
return th.zeros_like(param)
class MixedPrecisionTrainer:
def __init__(
self,
*,
model,
use_fp16=False,
fp16_scale_growth=1e-3,
initial_lg_loss_scale=INITIAL_LOG_LOSS_SCALE,
):
self.model = model
self.use_fp16 = use_fp16
self.fp16_scale_growth = fp16_scale_growth
self.model_params = list(self.model.parameters())
self.master_params = self.model_params
self.param_groups_and_shapes = None
self.lg_loss_scale = initial_lg_loss_scale
if self.use_fp16:
self.param_groups_and_shapes = get_param_groups_and_shapes(
self.model.named_parameters()
)
self.master_params = make_master_params(self.param_groups_and_shapes)
self.model.convert_to_fp16()
def zero_grad(self):
zero_grad(self.model_params)
def backward(self, loss: th.Tensor):
if self.use_fp16:
loss_scale = 2 ** self.lg_loss_scale
(loss * loss_scale).backward()
else:
loss.backward()
def optimize(self, opt: th.optim.Optimizer):
if self.use_fp16:
return self._optimize_fp16(opt)
else:
return self._optimize_normal(opt)
def _optimize_fp16(self, opt: th.optim.Optimizer):
logger.logkv_mean("lg_loss_scale", self.lg_loss_scale)
model_grads_to_master_grads(self.param_groups_and_shapes, self.master_params)
grad_norm, param_norm = self._compute_norms(grad_scale=2 ** self.lg_loss_scale)
if check_overflow(grad_norm):
self.lg_loss_scale -= 1
logger.log(f"Found NaN, decreased lg_loss_scale to {self.lg_loss_scale}")
zero_master_grads(self.master_params)
return False
logger.logkv_mean("grad_norm", grad_norm)
logger.logkv_mean("param_norm", param_norm)
self.master_params[0].grad.mul_(1.0 / (2 ** self.lg_loss_scale))
opt.step()
zero_master_grads(self.master_params)
master_params_to_model_params(self.param_groups_and_shapes, self.master_params)
self.lg_loss_scale += self.fp16_scale_growth
return True
def _optimize_normal(self, opt: th.optim.Optimizer):
grad_norm, param_norm = self._compute_norms()
logger.logkv_mean("grad_norm", grad_norm)
logger.logkv_mean("param_norm", param_norm)
opt.step()
return True
def _compute_norms(self, grad_scale=1.0):
grad_norm = 0.0
param_norm = 0.0
for p in self.master_params:
with th.no_grad():
param_norm += th.norm(p, p=2, dtype=th.float32).item() ** 2
if p.grad is not None:
grad_norm += th.norm(p.grad, p=2, dtype=th.float32).item() ** 2
return np.sqrt(grad_norm) / grad_scale, np.sqrt(param_norm)
def master_params_to_state_dict(self, master_params):
return master_params_to_state_dict(
self.model, self.param_groups_and_shapes, master_params, self.use_fp16
)
def state_dict_to_master_params(self, state_dict):
return state_dict_to_master_params(self.model, state_dict, self.use_fp16)
def check_overflow(value):
return (value == float("inf")) or (value == -float("inf")) or (value != value)
================================================
FILE: guided_diffusion/gaussian_diffusion.py
================================================
"""
This code started out as a PyTorch port of Ho et al's diffusion models:
https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py
Docstrings have been added, as well as DDIM sampling and a new collection of beta schedules.
"""
import enum
import math
import numpy as np
import torch as th
from torchvision import transforms
from .nn import mean_flat
from .losses import normal_kl, discretized_gaussian_log_likelihood
import os
from PIL import Image
def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
"""
Get a pre-defined beta schedule for the given name.
The beta schedule library consists of beta schedules which remain similar
in the limit of num_diffusion_timesteps.
Beta schedules may be added, but should not be removed or changed once
they are committed to maintain backwards compatibility.
"""
if schedule_name == "linear":
# Linear schedule from Ho et al, extended to work for any number of
# diffusion steps.
scale = 1000 / num_diffusion_timesteps
beta_start = scale * 0.0001
beta_end = scale * 0.02
return np.linspace(
beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64
)
elif schedule_name == "cosine":
return betas_for_alpha_bar(
num_diffusion_timesteps,
lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
)
else:
raise NotImplementedError(f"unknown beta schedule: {schedule_name}")
def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
"""
Create a beta schedule that discretizes the given alpha_t_bar function,
which defines the cumulative product of (1-beta) over time from t = [0,1].
:param num_diffusion_timesteps: the number of betas to produce.
:param alpha_bar: a lambda that takes an argument t from 0 to 1 and
produces the cumulative product of (1-beta) up to that
part of the diffusion process.
:param max_beta: the maximum beta to use; use values lower than 1 to
prevent singularities.
"""
betas = []
for i in range(num_diffusion_timesteps):
t1 = i / num_diffusion_timesteps
t2 = (i + 1) / num_diffusion_timesteps
betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
return np.array(betas)
class ModelMeanType(enum.Enum):
"""
Which type of output the model predicts.
"""
PREVIOUS_X = enum.auto() # the model predicts x_{t-1}
START_X = enum.auto() # the model predicts x_0
EPSILON = enum.auto() # the model predicts epsilon
class ModelVarType(enum.Enum):
"""
What is used as the model's output variance.
The LEARNED_RANGE option has been added to allow the model to predict
values between FIXED_SMALL and FIXED_LARGE, making its job easier.
"""
LEARNED = enum.auto()
FIXED_SMALL = enum.auto()
FIXED_LARGE = enum.auto()
LEARNED_RANGE = enum.auto()
class LossType(enum.Enum):
MSE = enum.auto() # use raw MSE loss (and KL when learning variances)
RESCALED_MSE = (
enum.auto()
) # use raw MSE loss (with RESCALED_KL when learning variances)
KL = enum.auto() # use the variational lower-bound
RESCALED_KL = enum.auto() # like KL, but rescale to estimate the full VLB
def is_vb(self):
return self == LossType.KL or self == LossType.RESCALED_KL
class GaussianDiffusion:
"""
Utilities for training and sampling diffusion models.
Ported directly from here, and then adapted over time to further experimentation.
https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42
:param betas: a 1-D numpy array of betas for each diffusion timestep,
starting at T and going to 1.
:param model_mean_type: a ModelMeanType determining what the model outputs.
:param model_var_type: a ModelVarType determining how variance is output.
:param loss_type: a LossType determining the loss function to use.
:param rescale_timesteps: if True, pass floating point timesteps into the
model so that they are always scaled like in the
original paper (0 to 1000).
"""
def __init__(
self,
*,
betas,
model_mean_type,
model_var_type,
loss_type,
rescale_timesteps=False,
):
self.model_mean_type = model_mean_type
self.model_var_type = model_var_type
self.loss_type = loss_type
self.rescale_timesteps = rescale_timesteps
# Use float64 for accuracy.
betas = np.array(betas, dtype=np.float64)
self.betas = betas
assert len(betas.shape) == 1, "betas must be 1-D"
assert (betas > 0).all() and (betas <= 1).all()
self.num_timesteps = int(betas.shape[0])
alphas = 1.0 - betas
self.alphas_cumprod = np.cumprod(alphas, axis=0)
self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)
# calculations for diffusion q(x_t | x_{t-1}) and others
self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)
# calculations for posterior q(x_{t-1} | x_t, x_0)
self.posterior_variance = (
betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
)
# log calculation clipped because the posterior variance is 0 at the
# beginning of the diffusion chain.
self.posterior_log_variance_clipped = np.log(
np.append(self.posterior_variance[1], self.posterior_variance[1:])
)
self.posterior_mean_coef1 = (
betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
)
self.posterior_mean_coef2 = (
(1.0 - self.alphas_cumprod_prev)
* np.sqrt(alphas)
/ (1.0 - self.alphas_cumprod)
)
def q_mean_variance(self, x_start, t):
"""
Get the distribution q(x_t | x_0).
:param x_start: the [N x C x ...] tensor of noiseless inputs.
:param t: the number of diffusion steps (minus 1). Here, 0 means one step.
:return: A tuple (mean, variance, log_variance), all of x_start's shape.
"""
mean = (
_extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
)
variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
log_variance = _extract_into_tensor(
self.log_one_minus_alphas_cumprod, t, x_start.shape
)
return mean, variance, log_variance
def q_sample(self, x_start, t, noise=None):
"""
Diffuse the data for a given number of diffusion steps.
In other words, sample from q(x_t | x_0).
:param x_start: the initial data batch.
:param t: the number of diffusion steps (minus 1). Here, 0 means one step.
:param noise: if specified, the split-out normal noise.
:return: A noisy version of x_start.
"""
raise NotImplementedError
if noise is None:
noise = th.randn_like(x_start)
assert noise.shape == x_start.shape
return (
_extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+ _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
* noise
)
def q_sample_blur(self, x_start, t, noise=None):
"""
Diffuse the data for a given number of diffusion steps.
In other words, sample from q(x_t | x_0).
:param x_start: the initial data batch.
:param t: the number of diffusion steps (minus 1). Here, 0 means one step.
:param noise: if specified, the split-out normal noise.
:return: A noisy version of x_start.
"""
assert th.all(t == t[0])
assert float((5*t[0]+1)/self.num_timesteps) >0 , f"t[0]: {t[0]}"
out = transforms.functional.gaussian_blur(x_start, 15, float((10*t[0]+1)/self.num_timesteps))
return out
def q_posterior_mean_variance(self, x_start, x_t, t):
"""
Compute the mean and variance of the diffusion posterior:
q(x_{t-1} | x_t, x_0)
"""
assert x_start.shape == x_t.shape
posterior_mean = (
_extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start
+ _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
)
posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape)
posterior_log_variance_clipped = _extract_into_tensor(
self.posterior_log_variance_clipped, t, x_t.shape
)
assert (
posterior_mean.shape[0]
== posterior_variance.shape[0]
== posterior_log_variance_clipped.shape[0]
== x_start.shape[0]
)
return posterior_mean, posterior_variance, posterior_log_variance_clipped
def p_mean_variance(
self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None
):
"""
Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
the initial x, x_0.
:param model: the model, which takes a signal and a batch of timesteps
as input.
:param x: the [N x C x ...] tensor at time t.
:param t: a 1-D Tensor of timesteps.
:param clip_denoised: if True, clip the denoised signal into [-1, 1].
:param denoised_fn: if not None, a function which applies to the
x_start prediction before it is used to sample. Applies before
clip_denoised.
:param model_kwargs: if not None, a dict of extra keyword arguments to
pass to the model. This can be used for conditioning.
:return: a dict with the following keys:
- 'mean': the model mean output.
- 'variance': the model variance output.
- 'log_variance': the log of 'variance'.
- 'pred_xstart': the prediction for x_0.
"""
if model_kwargs is None:
model_kwargs = {}
B, C = x.shape[:2]
assert t.shape == (B,)
model_output = model(x, self._scale_timesteps(t), **model_kwargs)
if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:
assert model_output.shape == (B, C * 2, *x.shape[2:])
model_output, model_var_values = th.split(model_output, C, dim=1)
if self.model_var_type == ModelVarType.LEARNED:
model_log_variance = model_var_values
model_variance = th.exp(model_log_variance)
else:
min_log = _extract_into_tensor(
self.posterior_log_variance_clipped, t, x.shape
)
max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)
# The model_var_values is [-1, 1] for [min_var, max_var].
frac = (model_var_values + 1) / 2
model_log_variance = frac * max_log + (1 - frac) * min_log
model_variance = th.exp(model_log_variance)
else:
model_variance, model_log_variance = {
# for fixedlarge, we set the initial (log-)variance like so
# to get a better decoder log likelihood.
ModelVarType.FIXED_LARGE: (
np.append(self.posterior_variance[1], self.betas[1:]),
np.log(np.append(self.posterior_variance[1], self.betas[1:])),
),
ModelVarType.FIXED_SMALL: (
self.posterior_variance,
self.posterior_log_variance_clipped,
),
}[self.model_var_type]
model_variance = _extract_into_tensor(model_variance, t, x.shape)
model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape)
def process_xstart(x):
if denoised_fn is not None:
x = denoised_fn(x)
if clip_denoised:
return x.clamp(-1, 1)
return x
if self.model_mean_type == ModelMeanType.PREVIOUS_X:
pred_xstart = process_xstart(
self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output)
)
model_mean = model_output
elif self.model_mean_type in [ModelMeanType.START_X, ModelMeanType.EPSILON]:
if self.model_mean_type == ModelMeanType.START_X:
pred_xstart = process_xstart(model_output)
else:
pred_xstart = process_xstart(
self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)
)
model_mean, _, _ = self.q_posterior_mean_variance(
x_start=pred_xstart, x_t=x, t=t
)
else:
raise NotImplementedError(self.model_mean_type)
assert (
model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
)
return {
"mean": model_mean,
"variance": model_variance,
"log_variance": model_log_variance,
"pred_xstart": pred_xstart,
}
def _predict_xstart_from_eps(self, x_t, t, eps):
assert x_t.shape == eps.shape
return (
_extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
- _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps
)
def _predict_xstart_from_xprev(self, x_t, t, xprev):
assert x_t.shape == xprev.shape
return ( # (xprev - coef2*x_t) / coef1
_extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev
- _extract_into_tensor(
self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape
)
* x_t
)
def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
return (
_extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
- pred_xstart
) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
def _scale_timesteps(self, t):
if self.rescale_timesteps:
return t.float() * (1000.0 / self.num_timesteps)
return t
def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
"""
Compute the mean for the previous step, given a function cond_fn that
computes the gradient of a conditional log probability with respect to
x. In particular, cond_fn computes grad(log(p(y|x))), and we want to
condition on y.
This uses the conditioning strategy from Sohl-Dickstein et al. (2015).
"""
gradient = cond_fn(x, self._scale_timesteps(t), **model_kwargs)
new_mean = (
p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float()
)
return new_mean
def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
"""
Compute what the p_mean_variance output would have been, should the
model's score function be conditioned by cond_fn.
See condition_mean() for details on cond_fn.
Unlike condition_mean(), this instead uses the conditioning strategy
from Song et al (2020).
"""
alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"])
eps = eps - (1 - alpha_bar).sqrt() * cond_fn(
x, self._scale_timesteps(t), **model_kwargs
)
out = p_mean_var.copy()
out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps)
out["mean"], _, _ = self.q_posterior_mean_variance(
x_start=out["pred_xstart"], x_t=x, t=t
)
return out
def p_sample(
self,
model,
x,
t,
clip_denoised=True,
denoised_fn=None,
cond_fn=None,
model_kwargs=None,
):
raise NotImplementedError
"""
Sample x_{t-1} from the model at the given timestep.
:param model: the model to sample from.
:param x: the current tensor at x_{t-1}.
:param t: the value of t, starting at 0 for the first diffusion step.
:param clip_denoised: if True, clip the x_start prediction to [-1, 1].
:param denoised_fn: if not None, a function which applies to the
x_start prediction before it is used to sample.
:param cond_fn: if not None, this is a gradient function that acts
similarly to the model.
:param model_kwargs: if not None, a dict of extra keyword arguments to
pass to the model. This can be used for conditioning.
:return: a dict containing the following keys:
- 'sample': a random sample from the model.
- 'pred_xstart': a prediction of x_0.
"""
out = self.p_mean_variance(
model,
x,
t,
clip_denoised=clip_denoised,
denoised_fn=denoised_fn,
model_kwargs=model_kwargs,
)
noise = th.randn_like(x)
nonzero_mask = (
(t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
) # no noise when t == 0
if cond_fn is not None:
out["mean"] = self.condition_mean(
cond_fn, out, x, t, model_kwargs=model_kwargs
)
sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise
return {"sample": sample, "pred_xstart": out["pred_xstart"]}
def p_sample_blur(
self,
model,
x,
t,
clip_denoised=True,
denoised_fn=None,
cond_fn=None,
model_kwargs=None,
):
"""
Sample x_{t-1} from the model at the given timestep.
:param model: the model to sample from.
:param x: the current tensor at x_{t-1}.
:param t: the value of t, starting at 0 for the first diffusion step.
:param clip_denoised: if True, clip the x_start prediction to [-1, 1].
:param denoised_fn: if not None, a function which applies to the
x_start prediction before it is used to sample.
:param cond_fn: if not None, this is a gradient function that acts
similarly to the model.
:param model_kwargs: if not None, a dict of extra keyword arguments to
pass to the model. This can be used for conditioning.
:return: a dict containing the following keys:
- 'sample': a random sample from the model.
- 'pred_xstart': a prediction of x_0.
"""
# out = self.p_mean_variance(
# model,
# x,
# t,
# clip_denoised=clip_denoised,
# denoised_fn=denoised_fn,
# model_kwargs=model_kwargs,
# )
# noise = th.randn_like(x)
# nonzero_mask = (
# (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
# ) # no noise when t == 0
if cond_fn is not None:
raise NotImplementedError
out["mean"] = self.condition_mean(
cond_fn, out, x, t, model_kwargs=model_kwargs
)
model_output = model(x, self._scale_timesteps(t), **model_kwargs)
# sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise
return {"sample": model_output}
def p_sample_loop(
self,
model,
shape,
noise=None,
clip_denoised=True,
denoised_fn=None,
cond_fn=None,
model_kwargs=None,
device=None,
progress=False,
):
"""
Generate samples from the model.
:param model: the model module.
:param shape: the shape of the samples, (N, C, H, W).
:param noise: if specified, the noise from the encoder to sample.
Should be of the same shape as `shape`.
:param clip_denoised: if True, clip x_start predictions to [-1, 1].
:param denoised_fn: if not None, a function which applies to the
x_start prediction before it is used to sample.
:param cond_fn: if not None, this is a gradient function that acts
similarly to the model.
:param model_kwargs: if not None, a dict of extra keyword arguments to
pass to the model. This can be used for conditioning.
:param device: if specified, the device to create the samples on.
If not specified, use a model parameter's device.
:param progress: if True, show a tqdm progress bar.
:return: a non-differentiable batch of samples.
"""
final = None
i=0
for sample in self.p_sample_loop_progressive(
model,
shape,
noise=noise,
clip_denoised=clip_denoised,
denoised_fn=denoised_fn,
cond_fn=cond_fn,
model_kwargs=model_kwargs,
device=device,
progress=progress,
):
print("p_sample_loop", i)
i+=1
tensor_imsave(sample["sample"][0], "./debug", f"{self.num_timesteps-i}.png", prt=True)
final = sample
return final["sample"]
def p_sample_loop_progressive(
self,
model,
shape,
noise=None,
clip_denoised=True,
denoised_fn=None,
cond_fn=None,
model_kwargs=None,
device=None,
progress=False,
):
"""
Generate samples from the model and yield intermediate samples from
each timestep of diffusion.
Arguments are the same as p_sample_loop().
Returns a generator over dicts, where each dict is the return value of
p_sample().
"""
if device is None:
device = next(model.parameters()).device
assert isinstance(shape, (tuple, list))
if noise is not None:
img = noise
else:
img = th.randn(*shape, device=device)
T = th.tensor(1000, device=device).repeat(shape[0])
print("T:", T)
img = self.q_sample_blur(img, T)
indices = list(range(self.num_timesteps))[::-1]
if progress:
# Lazy import so that we don't depend on tqdm.
from tqdm.auto import tqdm
indices = tqdm(indices)
for i in indices:
t = th.tensor([i] * shape[0], device=device)
with th.no_grad():
out = self.p_sample_blur(
model,
img,
t,
clip_denoised=clip_denoised,
denoised_fn=denoised_fn,
cond_fn=cond_fn,
model_kwargs=model_kwargs,
)
yield out
img = out["sample"]
def ddim_sample(
self,
model,
x,
t,
clip_denoised=True,
denoised_fn=None,
cond_fn=None,
model_kwargs=None,
eta=0.0,
):
"""
Sample x_{t-1} from the model using DDIM.
Same usage as p_sample().
"""
out = self.p_mean_variance(
model,
x,
t,
clip_denoised=clip_denoised,
denoised_fn=denoised_fn,
model_kwargs=model_kwargs,
)
if cond_fn is not None:
out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs)
# Usually our model outputs epsilon, but we re-derive it
# in case we used x_start or x_prev prediction.
eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"])
alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
sigma = (
eta
* th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))
* th.sqrt(1 - alpha_bar / alpha_bar_prev)
)
# Equation 12.
noise = th.randn_like(x)
mean_pred = (
out["pred_xstart"] * th.sqrt(alpha_bar_prev)
+ th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps
)
nonzero_mask = (
(t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
) # no noise when t == 0
sample = mean_pred + nonzero_mask * sigma * noise
return {"sample": sample, "pred_xstart": out["pred_xstart"]}
def ddim_reverse_sample(
self,
model,
x,
t,
clip_denoised=True,
denoised_fn=None,
model_kwargs=None,
eta=0.0,
):
"""
Sample x_{t+1} from the model using DDIM reverse ODE.
"""
assert eta == 0.0, "Reverse ODE only for deterministic path"
out = self.p_mean_variance(
model,
x,
t,
clip_denoised=clip_denoised,
denoised_fn=denoised_fn,
model_kwargs=model_kwargs,
)
# Usually our model outputs epsilon, but we re-derive it
# in case we used x_start or x_prev prediction.
eps = (
_extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x
- out["pred_xstart"]
) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape)
alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)
# Equation 12. reversed
mean_pred = (
out["pred_xstart"] * th.sqrt(alpha_bar_next)
+ th.sqrt(1 - alpha_bar_next) * eps
)
return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]}
def ddim_sample_loop(
self,
model,
shape,
noise=None,
clip_denoised=True,
denoised_fn=None,
cond_fn=None,
model_kwargs=None,
device=None,
progress=False,
eta=0.0,
):
"""
Generate samples from the model using DDIM.
Same usage as p_sample_loop().
"""
final = None
for sample in self.ddim_sample_loop_progressive(
model,
shape,
noise=noise,
clip_denoised=clip_denoised,
denoised_fn=denoised_fn,
cond_fn=cond_fn,
model_kwargs=model_kwargs,
device=device,
progress=progress,
eta=eta,
):
final = sample
return final["sample"]
def ddim_sample_loop_progressive(
self,
model,
shape,
noise=None,
clip_denoised=True,
denoised_fn=None,
cond_fn=None,
model_kwargs=None,
device=None,
progress=False,
eta=0.0,
):
"""
Use DDIM to sample from the model and yield intermediate samples from
each timestep of DDIM.
Same usage as p_sample_loop_progressive().
"""
if device is None:
device = next(model.parameters()).device
assert isinstance(shape, (tuple, list))
if noise is not None:
img = noise
else:
img = th.randn(*shape, device=device)
indices = list(range(self.num_timesteps))[::-1]
if progress:
# Lazy import so that we don't depend on tqdm.
from tqdm.auto import tqdm
indices = tqdm(indices)
for i in indices:
t = th.tensor([i] * shape[0], device=device)
with th.no_grad():
out = self.ddim_sample(
model,
img,
t,
clip_denoised=clip_denoised,
denoised_fn=denoised_fn,
cond_fn=cond_fn,
model_kwargs=model_kwargs,
eta=eta,
)
yield out
img = out["sample"]
def _vb_terms_bpd(
self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None
):
"""
Get a term for the variational lower-bound.
The resulting units are bits (rather than nats, as one might expect).
This allows for comparison to other papers.
:return: a dict with the following keys:
- 'output': a shape [N] tensor of NLLs or KLs.
- 'pred_xstart': the x_0 predictions.
"""
true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
x_start=x_start, x_t=x_t, t=t
)
out = self.p_mean_variance(
model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs
)
kl = normal_kl(
true_mean, true_log_variance_clipped, out["mean"], out["log_variance"]
)
kl = mean_flat(kl) / np.log(2.0)
decoder_nll = -discretized_gaussian_log_likelihood(
x_start, means=out["mean"], log_scales=0.5 * out["log_variance"]
)
assert decoder_nll.shape == x_start.shape
decoder_nll = mean_flat(decoder_nll) / np.log(2.0)
# At the first timestep return the decoder NLL,
# otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
output = th.where((t == 0), decoder_nll, kl)
return {"output": output, "pred_xstart": out["pred_xstart"]}
def training_losses(self, model, x_start, t, model_kwargs=None, noise=None):
"""
Compute training losses for a single timestep.
:param model: the model to evaluate loss on.
:param x_start: the [N x C x ...] tensor of inputs.
:param t: a batch of timestep indices.
:param model_kwargs: if not None, a dict of extra keyword arguments to
pass to the model. This can be used for conditioning.
:param noise: if specified, the specific Gaussian noise to try to remove.
:return: a dict with the key "loss" containing a tensor of shape [N].
Some mean or variance settings may also have other keys.
"""
if model_kwargs is None:
model_kwargs = {}
if noise is None:
noise = th.randn_like(x_start)
# x_t = self.q_sample(x_start, t, noise=noise)
x_t = self.q_sample_blur(x_start, t)
x_t_prev = self.q_sample_blur(x_start, t-1) # x_{t-1}
terms = {}
if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL:
raise NotImplementedError
terms["loss"] = self._vb_terms_bpd(
model=model,
x_start=x_start,
x_t=x_t,
t=t,
clip_denoised=False,
model_kwargs=model_kwargs,
)["output"]
if self.loss_type == LossType.RESCALED_KL:
terms["loss"] *= self.num_timesteps
elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE:
model_output = model(x_t, self._scale_timesteps(t), **model_kwargs)
if self.model_var_type in [
ModelVarType.LEARNED,
ModelVarType.LEARNED_RANGE,
]:
raise NotImplementedError
B, C = x_t.shape[:2]
assert model_output.shape == (B, C * 2, *x_t.shape[2:])
model_output, model_var_values = th.split(model_output, C, dim=1)
# Learn the variance using the variational bound, but don't let
# it affect our mean prediction.
frozen_out = th.cat([model_output.detach(), model_var_values], dim=1)
terms["vb"] = self._vb_terms_bpd(
model=lambda *args, r=frozen_out: r,
x_start=x_start,
x_t=x_t,
t=t,
clip_denoised=False,
)["output"]
if self.loss_type == LossType.RESCALED_MSE:
# Divide by 1000 for equivalence with initial implementation.
# Without a factor of 1/1000, the VB term hurts the MSE term.
terms["vb"] *= self.num_timesteps / 1000.0
# target = {
# ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance(
# x_start=x_start, x_t=x_t, t=t
# )[0],
# ModelMeanType.START_X: x_start,
# ModelMeanType.EPSILON: noise,
# }[self.model_mean_type]
target = x_t_prev
assert model_output.shape == target.shape == x_start.shape
terms["mse"] = mean_flat((target - model_output) ** 2)
if "vb" in terms:
terms["loss"] = terms["mse"] + terms["vb"]
else:
terms["loss"] = terms["mse"]
else:
raise NotImplementedError(self.loss_type)
return terms
def _prior_bpd(self, x_start):
"""
Get the prior KL term for the variational lower-bound, measured in
bits-per-dim.
This term can't be optimized, as it only depends on the encoder.
:param x_start: the [N x C x ...] tensor of inputs.
:return: a batch of [N] KL values (in bits), one per batch element.
"""
batch_size = x_start.shape[0]
t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
kl_prior = normal_kl(
mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0
)
return mean_flat(kl_prior) / np.log(2.0)
def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None):
"""
Compute the entire variational lower-bound, measured in bits-per-dim,
as well as other related quantities.
:param model: the model to evaluate loss on.
:param x_start: the [N x C x ...] tensor of inputs.
:param clip_denoised: if True, clip denoised samples.
:param model_kwargs: if not None, a dict of extra keyword arguments to
pass to the model. This can be used for conditioning.
:return: a dict containing the following keys:
- total_bpd: the total variational lower-bound, per batch element.
- prior_bpd: the prior term in the lower-bound.
- vb: an [N x T] tensor of terms in the lower-bound.
- xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep.
- mse: an [N x T] tensor of epsilon MSEs for each timestep.
"""
device = x_start.device
batch_size = x_start.shape[0]
xstart_mse = []
mse = []
for t in list(range(self.num_timesteps))[::-1]:
t_batch = th.tensor([t] * batch_size, device=device)
x_t = self.q_sample_blur(x_start=x_start, t=t_batch)
x_t_prev = self.q_sample_blur(x_start=x_start, t=t_batch-1)
# Calculate VLB term at the current timestep
with th.no_grad():
model_output = model(x_t, self._scale_timesteps(t), **model_kwargs)
mse.append(mean_flat((x_t_prev - model_output)**2))
mse = th.stack(mse, dim=1)
prior_bpd = self._prior_bpd(x_start)
return {'mse' : mse}
# device = x_start.device
# batch_size = x_start.shape[0]
# vb = []
# xstart_mse = []
# mse = []
# for t in list(range(self.num_timesteps))[::-1]:
# t_batch = th.tensor([t] * batch_size, device=device)
# noise = th.randn_like(x_start)
# x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise)
# # Calculate VLB term at the current timestep
# with th.no_grad():
# out = self._vb_terms_bpd(
# model,
# x_start=x_start,
# x_t=x_t,
# t=t_batch,
# clip_denoised=clip_denoised,
# model_kwargs=model_kwargs,
# )
# vb.append(out["output"])
# xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2))
# eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"])
# mse.append(mean_flat((eps - noise) ** 2))
# vb = th.stack(vb, dim=1)
# xstart_mse = th.stack(xstart_mse, dim=1)
# mse = th.stack(mse, dim=1)
# prior_bpd = self._prior_bpd(x_start)
# total_bpd = vb.sum(dim=1) + prior_bpd
# return {
# "total_bpd": total_bpd,
# "prior_bpd": prior_bpd,
# "vb": vb,
# "xstart_mse": xstart_mse,
# "mse": mse,
# }
def _extract_into_tensor(arr, timesteps, broadcast_shape):
"""
Extract values from a 1-D numpy array for a batch of indices.
:param arr: the 1-D numpy array.
:param timesteps: a tensor of indices into the array to extract.
:param broadcast_shape: a larger shape of K dimensions with the batch
dimension equal to the length of timesteps.
:return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
"""
res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
while len(res.shape) < len(broadcast_shape):
res = res[..., None]
return res.expand(broadcast_shape)
def tensor2pil(t):
img = transforms.functional.to_pil_image(denorm(t))
return img
def denorm(t):
return t*0.5 + 0.5
def tensor_imsave(t, path, fname, denormalization=True, prt=False):
# Save tensor as .png file
check_folder(path)
if denormalization: t = denorm(t)
t = th.clamp(input=t, min=0, max=1)
img = transforms.functional.to_pil_image(t.detach().cpu())
img.save(os.path.join(path,fname))
if prt:
print(f"Saved to {os.path.join(path,fname)}")
def pil_loader(path):
# open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
with open(path, 'rb') as f:##Automatic closing using with. rb = read binary
img = Image.open(f)
return img.convert('RGB')
def check_folder(log_dir):
if not os.path.exists(log_dir):
os.makedirs(log_dir)
return log_dir
================================================
FILE: guided_diffusion/image_datasets.py
================================================
import math
import random
from PIL import Image
import blobfile as bf
from mpi4py import MPI
import numpy as np
from torch.utils.data import DataLoader, Dataset
def load_data(
*,
data_dir,
batch_size,
image_size,
class_cond=False,
deterministic=False,
random_crop=False,
random_flip=True,
):
"""
For a dataset, create a generator over (images, kwargs) pairs.
Each images is an NCHW float tensor, and the kwargs dict contains zero or
more keys, each of which map to a batched Tensor of their own.
The kwargs dict can be used for class labels, in which case the key is "y"
and the values are integer tensors of class labels.
:param data_dir: a dataset directory.
:param batch_size: the batch size of each returned pair.
:param image_size: the size to which images are resized.
:param class_cond: if True, include a "y" key in returned dicts for class
label. If classes are not available and this is true, an
exception will be raised.
:param deterministic: if True, yield results in a deterministic order.
:param random_crop: if True, randomly crop the images for augmentation.
:param random_flip: if True, randomly flip the images for augmentation.
"""
if not data_dir:
raise ValueError("unspecified data directory")
all_files = _list_image_files_recursively(data_dir)
classes = None
if class_cond:
# Assume classes are the first part of the filename,
# before an underscore.
class_names = [bf.basename(path).split("_")[0] for path in all_files]
sorted_classes = {x: i for i, x in enumerate(sorted(set(class_names)))}
classes = [sorted_classes[x] for x in class_names]
dataset = ImageDataset(
image_size,
all_files,
classes=classes,
shard=MPI.COMM_WORLD.Get_rank(),
num_shards=MPI.COMM_WORLD.Get_size(),
random_crop=random_crop,
random_flip=random_flip,
)
if deterministic:
loader = DataLoader(
dataset, batch_size=batch_size, shuffle=False, num_workers=1, drop_last=True
)
else:
loader = DataLoader(
dataset, batch_size=batch_size, shuffle=True, num_workers=1, drop_last=True
)
while True:
yield from loader
def _list_image_files_recursively(data_dir):
results = []
for entry in sorted(bf.listdir(data_dir)):
full_path = bf.join(data_dir, entry)
ext = entry.split(".")[-1]
if "." in entry and ext.lower() in ["jpg", "jpeg", "png", "gif"]:
results.append(full_path)
elif bf.isdir(full_path):
results.extend(_list_image_files_recursively(full_path))
return results
class ImageDataset(Dataset):
def __init__(
self,
resolution,
image_paths,
classes=None,
shard=0,
num_shards=1,
random_crop=False,
random_flip=True,
):
super().__init__()
self.resolution = resolution
self.local_images = image_paths[shard:][::num_shards]
self.local_classes = None if classes is None else classes[shard:][::num_shards]
self.random_crop = random_crop
self.random_flip = random_flip
def __len__(self):
return len(self.local_images)
def __getitem__(self, idx):
path = self.local_images[idx]
with bf.BlobFile(path, "rb") as f:
pil_image = Image.open(f)
pil_image.load()
pil_image = pil_image.convert("RGB")
if self.random_crop:
arr = random_crop_arr(pil_image, self.resolution)
else:
arr = center_crop_arr(pil_image, self.resolution)
if self.random_flip and random.random() < 0.5:
arr = arr[:, ::-1]
arr = arr.astype(np.float32) / 127.5 - 1
out_dict = {}
if self.local_classes is not None:
out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64)
return np.transpose(arr, [2, 0, 1]), out_dict
def center_crop_arr(pil_image, image_size):
# We are not on a new enough PIL to support the `reducing_gap`
# argument, which uses BOX downsampling at powers of two first.
# Thus, we do it by hand to improve downsample quality.
while min(*pil_image.size) >= 2 * image_size:
pil_image = pil_image.resize(
tuple(x // 2 for x in pil_image.size), resample=Image.BOX
)
scale = image_size / min(*pil_image.size)
pil_image = pil_image.resize(
tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
)
arr = np.array(pil_image)
crop_y = (arr.shape[0] - image_size) // 2
crop_x = (arr.shape[1] - image_size) // 2
return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size]
def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0):
min_smaller_dim_size = math.ceil(image_size / max_crop_frac)
max_smaller_dim_size = math.ceil(image_size / min_crop_frac)
smaller_dim_size = random.randrange(min_smaller_dim_size, max_smaller_dim_size + 1)
# We are not on a new enough PIL to support the `reducing_gap`
# argument, which uses BOX downsampling at powers of two first.
# Thus, we do it by hand to improve downsample quality.
while min(*pil_image.size) >= 2 * smaller_dim_size:
pil_image = pil_image.resize(
tuple(x // 2 for x in pil_image.size), resample=Image.BOX
)
scale = smaller_dim_size / min(*pil_image.size)
pil_image = pil_image.resize(
tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
)
arr = np.array(pil_image)
crop_y = random.randrange(arr.shape[0] - image_size + 1)
crop_x = random.randrange(arr.shape[1] - image_size + 1)
return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size]
================================================
FILE: guided_diffusion/logger.py
================================================
"""
Logger copied from OpenAI baselines to avoid extra RL-based dependencies:
https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/logger.py
"""
import os
import sys
import shutil
import os.path as osp
import json
import time
import datetime
import tempfile
import warnings
from collections import defaultdict
from contextlib import contextmanager
DEBUG = 10
INFO = 20
WARN = 30
ERROR = 40
DISABLED = 50
class KVWriter(object):
def writekvs(self, kvs):
raise NotImplementedError
class SeqWriter(object):
def writeseq(self, seq):
raise NotImplementedError
class HumanOutputFormat(KVWriter, SeqWriter):
def __init__(self, filename_or_file):
if isinstance(filename_or_file, str):
self.file = open(filename_or_file, "wt")
self.own_file = True
else:
assert hasattr(filename_or_file, "read"), (
"expected file or str, got %s" % filename_or_file
)
self.file = filename_or_file
self.own_file = False
def writekvs(self, kvs):
# Create strings for printing
key2str = {}
for (key, val) in sorted(kvs.items()):
if hasattr(val, "__float__"):
valstr = "%-8.3g" % val
else:
valstr = str(val)
key2str[self._truncate(key)] = self._truncate(valstr)
# Find max widths
if len(key2str) == 0:
print("WARNING: tried to write empty key-value dict")
return
else:
keywidth = max(map(len, key2str.keys()))
valwidth = max(map(len, key2str.values()))
# Write out the data
dashes = "-" * (keywidth + valwidth + 7)
lines = [dashes]
for (key, val) in sorted(key2str.items(), key=lambda kv: kv[0].lower()):
lines.append(
"| %s%s | %s%s |"
% (key, " " * (keywidth - len(key)), val, " " * (valwidth - len(val)))
)
lines.append(dashes)
self.file.write("\n".join(lines) + "\n")
# Flush the output to the file
self.file.flush()
def _truncate(self, s):
maxlen = 30
return s[: maxlen - 3] + "..." if len(s) > maxlen else s
def writeseq(self, seq):
seq = list(seq)
for (i, elem) in enumerate(seq):
self.file.write(elem)
if i < len(seq) - 1: # add space unless this is the last one
self.file.write(" ")
self.file.write("\n")
self.file.flush()
def close(self):
if self.own_file:
self.file.close()
class JSONOutputFormat(KVWriter):
def __init__(self, filename):
self.file = open(filename, "wt")
def writekvs(self, kvs):
for k, v in sorted(kvs.items()):
if hasattr(v, "dtype"):
kvs[k] = float(v)
self.file.write(json.dumps(kvs) + "\n")
self.file.flush()
def close(self):
self.file.close()
class CSVOutputFormat(KVWriter):
def __init__(self, filename):
self.file = open(filename, "w+t")
self.keys = []
self.sep = ","
def writekvs(self, kvs):
# Add our current row to the history
extra_keys = list(kvs.keys() - self.keys)
extra_keys.sort()
if extra_keys:
self.keys.extend(extra_keys)
self.file.seek(0)
lines = self.file.readlines()
self.file.seek(0)
for (i, k) in enumerate(self.keys):
if i > 0:
self.file.write(",")
self.file.write(k)
self.file.write("\n")
for line in lines[1:]:
self.file.write(line[:-1])
self.file.write(self.sep * len(extra_keys))
self.file.write("\n")
for (i, k) in enumerate(self.keys):
if i > 0:
self.file.write(",")
v = kvs.get(k)
if v is not None:
self.file.write(str(v))
self.file.write("\n")
self.file.flush()
def close(self):
self.file.close()
class TensorBoardOutputFormat(KVWriter):
"""
Dumps key/value pairs into TensorBoard's numeric format.
"""
def __init__(self, dir):
os.makedirs(dir, exist_ok=True)
self.dir = dir
self.step = 1
prefix = "events"
path = osp.join(osp.abspath(dir), prefix)
import tensorflow as tf
from tensorflow.python import pywrap_tensorflow
from tensorflow.core.util import event_pb2
from tensorflow.python.util import compat
self.tf = tf
self.event_pb2 = event_pb2
self.pywrap_tensorflow = pywrap_tensorflow
self.writer = pywrap_tensorflow.EventsWriter(compat.as_bytes(path))
def writekvs(self, kvs):
def summary_val(k, v):
kwargs = {"tag": k, "simple_value": float(v)}
return self.tf.Summary.Value(**kwargs)
summary = self.tf.Summary(value=[summary_val(k, v) for k, v in kvs.items()])
event = self.event_pb2.Event(wall_time=time.time(), summary=summary)
event.step = (
self.step
) # is there any reason why you'd want to specify the step?
self.writer.WriteEvent(event)
self.writer.Flush()
self.step += 1
def close(self):
if self.writer:
self.writer.Close()
self.writer = None
def make_output_format(format, ev_dir, log_suffix=""):
os.makedirs(ev_dir, exist_ok=True)
if format == "stdout":
return HumanOutputFormat(sys.stdout)
elif format == "log":
return HumanOutputFormat(osp.join(ev_dir, "log%s.txt" % log_suffix))
elif format == "json":
return JSONOutputFormat(osp.join(ev_dir, "progress%s.json" % log_suffix))
elif format == "csv":
return CSVOutputFormat(osp.join(ev_dir, "progress%s.csv" % log_suffix))
elif format == "tensorboard":
return TensorBoardOutputFormat(osp.join(ev_dir, "tb%s" % log_suffix))
else:
raise ValueError("Unknown format specified: %s" % (format,))
# ================================================================
# API
# ================================================================
def logkv(key, val):
"""
Log a value of some diagnostic
Call this once for each diagnostic quantity, each iteration
If called many times, last value will be used.
"""
get_current().logkv(key, val)
def logkv_mean(key, val):
"""
The same as logkv(), but if called many times, values averaged.
"""
get_current().logkv_mean(key, val)
def logkvs(d):
"""
Log a dictionary of key-value pairs
"""
for (k, v) in d.items():
logkv(k, v)
def dumpkvs():
"""
Write all of the diagnostics from the current iteration
"""
return get_current().dumpkvs()
def getkvs():
return get_current().name2val
def log(*args, level=INFO):
"""
Write the sequence of args, with no separators, to the console and output files (if you've configured an output file).
"""
get_current().log(*args, level=level)
def debug(*args):
log(*args, level=DEBUG)
def info(*args):
log(*args, level=INFO)
def warn(*args):
log(*args, level=WARN)
def error(*args):
log(*args, level=ERROR)
def set_level(level):
"""
Set logging threshold on current logger.
"""
get_current().set_level(level)
def set_comm(comm):
get_current().set_comm(comm)
def get_dir():
"""
Get directory that log files are being written to.
will be None if there is no output directory (i.e., if you didn't call start)
"""
return get_current().get_dir()
record_tabular = logkv
dump_tabular = dumpkvs
@contextmanager
def profile_kv(scopename):
logkey = "wait_" + scopename
tstart = time.time()
try:
yield
finally:
get_current().name2val[logkey] += time.time() - tstart
def profile(n):
"""
Usage:
@profile("my_func")
def my_func(): code
"""
def decorator_with_name(func):
def func_wrapper(*args, **kwargs):
with profile_kv(n):
return func(*args, **kwargs)
return func_wrapper
return decorator_with_name
# ================================================================
# Backend
# ================================================================
def get_current():
if Logger.CURRENT is None:
_configure_default_logger()
return Logger.CURRENT
class Logger(object):
DEFAULT = None # A logger with no output files. (See right below class definition)
# So that you can still log to the terminal without setting up any output files
CURRENT = None # Current logger being used by the free functions above
def __init__(self, dir, output_formats, comm=None):
self.name2val = defaultdict(float) # values this iteration
self.name2cnt = defaultdict(int)
self.level = INFO
self.dir = dir
self.output_formats = output_formats
self.comm = comm
# Logging API, forwarded
# ----------------------------------------
def logkv(self, key, val):
self.name2val[key] = val
def logkv_mean(self, key, val):
oldval, cnt = self.name2val[key], self.name2cnt[key]
self.name2val[key] = oldval * cnt / (cnt + 1) + val / (cnt + 1)
self.name2cnt[key] = cnt + 1
def dumpkvs(self):
if self.comm is None:
d = self.name2val
else:
d = mpi_weighted_mean(
self.comm,
{
name: (val, self.name2cnt.get(name, 1))
for (name, val) in self.name2val.items()
},
)
if self.comm.rank != 0:
d["dummy"] = 1 # so we don't get a warning about empty dict
out = d.copy() # Return the dict for unit testing purposes
for fmt in self.output_formats:
if isinstance(fmt, KVWriter):
fmt.writekvs(d)
self.name2val.clear()
self.name2cnt.clear()
return out
def log(self, *args, level=INFO):
if self.level <= level:
self._do_log(args)
# Configuration
# ----------------------------------------
def set_level(self, level):
self.level = level
def set_comm(self, comm):
self.comm = comm
def get_dir(self):
return self.dir
def close(self):
for fmt in self.output_formats:
fmt.close()
# Misc
# ----------------------------------------
def _do_log(self, args):
for fmt in self.output_formats:
if isinstance(fmt, SeqWriter):
fmt.writeseq(map(str, args))
def get_rank_without_mpi_import():
# check environment variables here instead of importing mpi4py
# to avoid calling MPI_Init() when this module is imported
for varname in ["PMI_RANK", "OMPI_COMM_WORLD_RANK"]:
if varname in os.environ:
return int(os.environ[varname])
return 0
def mpi_weighted_mean(comm, local_name2valcount):
"""
Copied from: https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/common/mpi_util.py#L110
Perform a weighted average over dicts that are each on a different node
Input: local_name2valcount: dict mapping key -> (value, count)
Returns: key -> mean
"""
all_name2valcount = comm.gather(local_name2valcount)
if comm.rank == 0:
name2sum = defaultdict(float)
name2count = defaultdict(float)
for n2vc in all_name2valcount:
for (name, (val, count)) in n2vc.items():
try:
val = float(val)
except ValueError:
if comm.rank == 0:
warnings.warn(
"WARNING: tried to compute mean on non-float {}={}".format(
name, val
)
)
else:
name2sum[name] += val * count
name2count[name] += count
return {name: name2sum[name] / name2count[name] for name in name2sum}
else:
return {}
def configure(dir=None, format_strs=None, comm=None, log_suffix=""):
"""
If comm is provided, average all numerical stats across that comm
"""
dir = "./cifar10-forwardblur-fromblurrednoise"
if dir is None:
dir = os.getenv("OPENAI_LOGDIR")
if dir is None:
dir = osp.join(
tempfile.gettempdir(),
datetime.datetime.now().strftime("openai-%Y-%m-%d-%H-%M-%S-%f"),
)
assert isinstance(dir, str)
dir = os.path.expanduser(dir)
os.makedirs(os.path.expanduser(dir), exist_ok=True)
rank = get_rank_without_mpi_import()
if rank > 0:
log_suffix = log_suffix + "-rank%03i" % rank
if format_strs is None:
if rank == 0:
format_strs = os.getenv("OPENAI_LOG_FORMAT", "stdout,log,csv").split(",")
else:
format_strs = os.getenv("OPENAI_LOG_FORMAT_MPI", "log").split(",")
format_strs = filter(None, format_strs)
output_formats = [make_output_format(f, dir, log_suffix) for f in format_strs]
Logger.CURRENT = Logger(dir=dir, output_formats=output_formats, comm=comm)
if output_formats:
log("Logging to %s" % dir)
def _configure_default_logger():
configure()
Logger.DEFAULT = Logger.CURRENT
def reset():
if Logger.CURRENT is not Logger.DEFAULT:
Logger.CURRENT.close()
Logger.CURRENT = Logger.DEFAULT
log("Reset logger")
@contextmanager
def scoped_configure(dir=None, format_strs=None, comm=None):
prevlogger = Logger.CURRENT
configure(dir=dir, format_strs=format_strs, comm=comm)
try:
yield
finally:
Logger.CURRENT.close()
Logger.CURRENT = prevlogger
================================================
FILE: guided_diffusion/losses.py
================================================
"""
Helpers for various likelihood-based losses. These are ported from the original
Ho et al. diffusion models codebase:
https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/utils.py
"""
import numpy as np
import torch as th
def normal_kl(mean1, logvar1, mean2, logvar2):
"""
Compute the KL divergence between two gaussians.
Shapes are automatically broadcasted, so batches can be compared to
scalars, among other use cases.
"""
tensor = None
for obj in (mean1, logvar1, mean2, logvar2):
if isinstance(obj, th.Tensor):
tensor = obj
break
assert tensor is not None, "at least one argument must be a Tensor"
# Force variances to be Tensors. Broadcasting helps convert scalars to
# Tensors, but it does not work for th.exp().
logvar1, logvar2 = [
x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor)
for x in (logvar1, logvar2)
]
return 0.5 * (
-1.0
+ logvar2
- logvar1
+ th.exp(logvar1 - logvar2)
+ ((mean1 - mean2) ** 2) * th.exp(-logvar2)
)
def approx_standard_normal_cdf(x):
"""
A fast approximation of the cumulative distribution function of the
standard normal.
"""
return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3))))
def discretized_gaussian_log_likelihood(x, *, means, log_scales):
"""
Compute the log-likelihood of a Gaussian distribution discretizing to a
given image.
:param x: the target images. It is assumed that this was uint8 values,
rescaled to the range [-1, 1].
:param means: the Gaussian mean Tensor.
:param log_scales: the Gaussian log stddev Tensor.
:return: a tensor like x of log probabilities (in nats).
"""
assert x.shape == means.shape == log_scales.shape
centered_x = x - means
inv_stdv = th.exp(-log_scales)
plus_in = inv_stdv * (centered_x + 1.0 / 255.0)
cdf_plus = approx_standard_normal_cdf(plus_in)
min_in = inv_stdv * (centered_x - 1.0 / 255.0)
cdf_min = approx_standard_normal_cdf(min_in)
log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12))
log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12))
cdf_delta = cdf_plus - cdf_min
log_probs = th.where(
x < -0.999,
log_cdf_plus,
th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))),
)
assert log_probs.shape == x.shape
return log_probs
================================================
FILE: guided_diffusion/nn.py
================================================
"""
Various utilities for neural networks.
"""
import math
import torch as th
import torch.nn as nn
# PyTorch 1.7 has SiLU, but we support PyTorch 1.5.
class SiLU(nn.Module):
def forward(self, x):
return x * th.sigmoid(x)
class GroupNorm32(nn.GroupNorm):
def forward(self, x):
return super().forward(x.float()).type(x.dtype)
def conv_nd(dims, *args, **kwargs):
"""
Create a 1D, 2D, or 3D convolution module.
"""
if dims == 1:
return nn.Conv1d(*args, **kwargs)
elif dims == 2:
return nn.Conv2d(*args, **kwargs)
elif dims == 3:
return nn.Conv3d(*args, **kwargs)
raise ValueError(f"unsupported dimensions: {dims}")
def linear(*args, **kwargs):
"""
Create a linear module.
"""
return nn.Linear(*args, **kwargs)
def avg_pool_nd(dims, *args, **kwargs):
"""
Create a 1D, 2D, or 3D average pooling module.
"""
if dims == 1:
return nn.AvgPool1d(*args, **kwargs)
elif dims == 2:
return nn.AvgPool2d(*args, **kwargs)
elif dims == 3:
return nn.AvgPool3d(*args, **kwargs)
raise ValueError(f"unsupported dimensions: {dims}")
def update_ema(target_params, source_params, rate=0.99):
"""
Update target parameters to be closer to those of source parameters using
an exponential moving average.
:param target_params: the target parameter sequence.
:param source_params: the source parameter sequence.
:param rate: the EMA rate (closer to 1 means slower).
"""
for targ, src in zip(target_params, source_params):
targ.detach().mul_(rate).add_(src, alpha=1 - rate)
def zero_module(module):
"""
Zero out the parameters of a module and return it.
"""
for p in module.parameters():
p.detach().zero_()
return module
def scale_module(module, scale):
"""
Scale the parameters of a module and return it.
"""
for p in module.parameters():
p.detach().mul_(scale)
return module
def mean_flat(tensor):
"""
Take the mean over all non-batch dimensions.
"""
return tensor.mean(dim=list(range(1, len(tensor.shape))))
def normalization(channels):
"""
Make a standard normalization layer.
:param channels: number of input channels.
:return: an nn.Module for normalization.
"""
return GroupNorm32(32, channels)
def timestep_embedding(timesteps, dim, max_period=10000):
"""
Create sinusoidal timestep embeddings.
:param timesteps: a 1-D Tensor of N indices, one per batch element.
These may be fractional.
:param dim: the dimension of the output.
:param max_period: controls the minimum frequency of the embeddings.
:return: an [N x dim] Tensor of positional embeddings.
"""
half = dim // 2
freqs = th.exp(
-math.log(max_period) * th.arange(start=0, end=half, dtype=th.float32) / half
).to(device=timesteps.device)
args = timesteps[:, None].float() * freqs[None]
embedding = th.cat([th.cos(args), th.sin(args)], dim=-1)
if dim % 2:
embedding = th.cat([embedding, th.zeros_like(embedding[:, :1])], dim=-1)
return embedding
def checkpoint(func, inputs, params, flag):
"""
Evaluate a function without caching intermediate activations, allowing for
reduced memory at the expense of extra compute in the backward pass.
:param func: the function to evaluate.
:param inputs: the argument sequence to pass to `func`.
:param params: a sequence of parameters `func` depends on but does not
explicitly take as arguments.
:param flag: if False, disable gradient checkpointing.
"""
if flag:
args = tuple(inputs) + tuple(params)
return CheckpointFunction.apply(func, len(inputs), *args)
else:
return func(*inputs)
class CheckpointFunction(th.autograd.Function):
@staticmethod
def forward(ctx, run_function, length, *args):
ctx.run_function = run_function
ctx.input_tensors = list(args[:length])
ctx.input_params = list(args[length:])
with th.no_grad():
output_tensors = ctx.run_function(*ctx.input_tensors)
return output_tensors
@staticmethod
def backward(ctx, *output_grads):
ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
with th.enable_grad():
# Fixes a bug where the first op in run_function modifies the
# Tensor storage in place, which is not allowed for detach()'d
# Tensors.
shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
output_tensors = ctx.run_function(*shallow_copies)
input_grads = th.autograd.grad(
output_tensors,
ctx.input_tensors + ctx.input_params,
output_grads,
allow_unused=True,
)
del ctx.input_tensors
del ctx.input_params
del output_tensors
return (None, None) + input_grads
================================================
FILE: guided_diffusion/resample.py
================================================
from abc import ABC, abstractmethod
import numpy as np
import torch as th
import torch.distributed as dist
def create_named_schedule_sampler(name, diffusion):
"""
Create a ScheduleSampler from a library of pre-defined samplers.
:param name: the name of the sampler.
:param diffusion: the diffusion object to sample for.
"""
if name == "uniform":
return UniformSampler(diffusion)
elif name == "loss-second-moment":
return LossSecondMomentResampler(diffusion)
else:
raise NotImplementedError(f"unknown schedule sampler: {name}")
class ScheduleSampler(ABC):
"""
A distribution over timesteps in the diffusion process, intended to reduce
variance of the objective.
By default, samplers perform unbiased importance sampling, in which the
objective's mean is unchanged.
However, subclasses may override sample() to change how the resampled
terms are reweighted, allowing for actual changes in the objective.
"""
@abstractmethod
def weights(self):
"""
Get a numpy array of weights, one per diffusion step.
The weights needn't be normalized, but must be positive.
"""
def sample(self, batch_size, device):
"""
Importance-sample timesteps for a batch.
:param batch_size: the number of timesteps.
:param device: the torch device to save to.
:return: a tuple (timesteps, weights):
- timesteps: a tensor of timestep indices.
- weights: a tensor of weights to scale the resulting losses.
"""
w = self.weights()
p = w / np.sum(w) # loss별 weight를 normalize
indices_np = np.random.choice(len(p), size=(batch_size,), p=p)[0]
if indices_np == 0: indices_np += 1
assert indices_np > 0, f"indices_np : {indices_np}"
indices_np = np.repeat(indices_np, batch_size)
indices = th.from_numpy(indices_np).long().to(device)
weights_np = 1 / (len(p) * p[indices_np])
weights = th.from_numpy(weights_np).float().to(device)
return indices, weights
class UniformSampler(ScheduleSampler):
def __init__(self, diffusion):
self.diffusion = diffusion
self._weights = np.ones([diffusion.num_timesteps])
def weights(self):
return self._weights
class LossAwareSampler(ScheduleSampler):
def update_with_local_losses(self, local_ts, local_losses):
"""
Update the reweighting using losses from a model.
Call this method from each rank with a batch of timesteps and the
corresponding losses for each of those timesteps.
This method will perform synchronization to make sure all of the ranks
maintain the exact same reweighting.
:param local_ts: an integer Tensor of timesteps.
:param local_losses: a 1D Tensor of losses.
"""
batch_sizes = [
th.tensor([0], dtype=th.int32, device=local_ts.device)
for _ in range(dist.get_world_size())
]
dist.all_gather(
batch_sizes,
th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device),
)
# Pad all_gather batches to be the maximum batch size.
batch_sizes = [x.item() for x in batch_sizes]
max_bs = max(batch_sizes)
timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes]
loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes]
dist.all_gather(timestep_batches, local_ts)
dist.all_gather(loss_batches, local_losses)
timesteps = [
x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs]
]
losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]]
self.update_with_all_losses(timesteps, losses)
@abstractmethod
def update_with_all_losses(self, ts, losses):
"""
Update the reweighting using losses from a model.
Sub-classes should override this method to update the reweighting
using losses from the model.
This method directly updates the reweighting without synchronizing
between workers. It is called by update_with_local_losses from all
ranks with identical arguments. Thus, it should have deterministic
behavior to maintain state across workers.
:param ts: a list of int timesteps.
:param losses: a list of float losses, one per timestep.
"""
class LossSecondMomentResampler(LossAwareSampler):
def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001):
raise NotImplementedError
self.diffusion = diffusion
self.history_per_term = history_per_term
self.uniform_prob = uniform_prob
self._loss_history = np.zeros(
[diffusion.num_timesteps, history_per_term], dtype=np.float64
)
self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int)
def weights(self):
if not self._warmed_up():
return np.ones([self.diffusion.num_timesteps], dtype=np.float64)
weights = np.sqrt(np.mean(self._loss_history ** 2, axis=-1))
weights /= np.sum(weights)
weights *= 1 - self.uniform_prob
weights += self.uniform_prob / len(weights)
return weights
def update_with_all_losses(self, ts, losses):
for t, loss in zip(ts, losses):
if self._loss_counts[t] == self.history_per_term:
# Shift out the oldest loss term.
self._loss_history[t, :-1] = self._loss_history[t, 1:]
self._loss_history[t, -1] = loss
else:
self._loss_history[t, self._loss_counts[t]] = loss
self._loss_counts[t] += 1
def _warmed_up(self):
return (self._loss_counts == self.history_per_term).all()
================================================
FILE: guided_diffusion/respace.py
================================================
import numpy as np
import torch as th
from .gaussian_diffusion import GaussianDiffusion
def space_timesteps(num_timesteps, section_counts):
"""
Create a list of timesteps to use from an original diffusion process,
given the number of timesteps we want to take from equally-sized portions
of the original process.
For example, if there's 300 timesteps and the section counts are [10,15,20]
then the first 100 timesteps are strided to be 10 timesteps, the second 100
are strided to be 15 timesteps, and the final 100 are strided to be 20.
If the stride is a string starting with "ddim", then the fixed striding
from the DDIM paper is used, and only one section is allowed.
:param num_timesteps: the number of diffusion steps in the original
process to divide up.
:param section_counts: either a list of numbers, or a string containing
comma-separated numbers, indicating the step count
per section. As a special case, use "ddimN" where N
is a number of steps to use the striding from the
DDIM paper.
:return: a set of diffusion steps from the original process to use.
"""
if isinstance(section_counts, str):
if section_counts.startswith("ddim"):
desired_count = int(section_counts[len("ddim") :])
for i in range(1, num_timesteps):
if len(range(0, num_timesteps, i)) == desired_count:
return set(range(0, num_timesteps, i))
raise ValueError(
f"cannot create exactly {num_timesteps} steps with an integer stride"
)
section_counts = [int(x) for x in section_counts.split(",")]
size_per = num_timesteps // len(section_counts)
extra = num_timesteps % len(section_counts)
start_idx = 0
all_steps = []
for i, section_count in enumerate(section_counts):
size = size_per + (1 if i < extra else 0)
if size < section_count:
raise ValueError(
f"cannot divide section of {size} steps into {section_count}"
)
if section_count <= 1:
frac_stride = 1
else:
frac_stride = (size - 1) / (section_count - 1)
cur_idx = 0.0
taken_steps = []
for _ in range(section_count):
taken_steps.append(start_idx + round(cur_idx))
cur_idx += frac_stride
all_steps += taken_steps
start_idx += size
return set(all_steps)
class SpacedDiffusion(GaussianDiffusion):
"""
A diffusion process which can skip steps in a base diffusion process.
:param use_timesteps: a collection (sequence or set) of timesteps from the
original diffusion process to retain.
:param kwargs: the kwargs to create the base diffusion process.
"""
def __init__(self, use_timesteps, **kwargs):
self.use_timesteps = set(use_timesteps)
self.timestep_map = []
self.original_num_steps = len(kwargs["betas"])
base_diffusion = GaussianDiffusion(**kwargs) # pylint: disable=missing-kwoa
last_alpha_cumprod = 1.0
new_betas = []
for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod):
if i in self.use_timesteps:
new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)
last_alpha_cumprod = alpha_cumprod
self.timestep_map.append(i)
kwargs["betas"] = np.array(new_betas)
super().__init__(**kwargs)
def p_mean_variance(
self, model, *args, **kwargs
): # pylint: disable=signature-differs
return super().p_mean_variance(self._wrap_model(model), *args, **kwargs)
def training_losses(
self, model, *args, **kwargs
): # pylint: disable=signature-differs
return super().training_losses(self._wrap_model(model), *args, **kwargs)
def condition_mean(self, cond_fn, *args, **kwargs):
return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs)
def condition_score(self, cond_fn, *args, **kwargs):
return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs)
def _wrap_model(self, model):
if isinstance(model, _WrappedModel):
return model
return _WrappedModel(
model, self.timestep_map, self.rescale_timesteps, self.original_num_steps
)
def _scale_timesteps(self, t):
# Scaling is done by the wrapped model.
return t
class _WrappedModel:
def __init__(self, model, timestep_map, rescale_timesteps, original_num_steps):
self.model = model
self.timestep_map = timestep_map
self.rescale_timesteps = rescale_timesteps
self.original_num_steps = original_num_steps
def __call__(self, x, ts, **kwargs):
map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)
new_ts = map_tensor[ts]
if self.rescale_timesteps:
new_ts = new_ts.float() * (1000.0 / self.original_num_steps)
return self.model(x, new_ts, **kwargs)
================================================
FILE: guided_diffusion/script_util.py
================================================
import argparse
import inspect
from . import gaussian_diffusion as gd
from .respace import SpacedDiffusion, space_timesteps
from .unet import SuperResModel, UNetModel, EncoderUNetModel
NUM_CLASSES = 1000
def diffusion_defaults():
"""
Defaults for image and classifier training.
"""
return dict(
learn_sigma=False,
diffusion_steps=1000,
noise_schedule="linear",
timestep_respacing="",
use_kl=False,
predict_xstart=False,
rescale_timesteps=False,
rescale_learned_sigmas=False,
)
def classifier_defaults():
"""
Defaults for classifier models.
"""
return dict(
image_size=64,
classifier_use_fp16=False,
classifier_width=128,
classifier_depth=2,
classifier_attention_resolutions="32,16,8", # 16
classifier_use_scale_shift_norm=True, # False
classifier_resblock_updown=True, # False
classifier_pool="attention",
)
def model_and_diffusion_defaults():
"""
Defaults for image training.
"""
res = dict(
image_size=64,
num_channels=128,
num_res_blocks=2,
num_heads=4,
num_heads_upsample=-1,
num_head_channels=-1,
attention_resolutions="16,8",
channel_mult="",
dropout=0.0,
class_cond=False,
use_checkpoint=False,
use_scale_shift_norm=True,
resblock_updown=False,
use_fp16=False,
use_new_attention_order=False,
)
res.update(diffusion_defaults())
return res
def classifier_and_diffusion_defaults():
res = classifier_defaults()
res.update(diffusion_defaults())
return res
def create_model_and_diffusion(
image_size,
class_cond,
learn_sigma,
num_channels,
num_res_blocks,
channel_mult,
num_heads,
num_head_channels,
num_heads_upsample,
attention_resolutions,
dropout,
diffusion_steps,
noise_schedule,
timestep_respacing,
use_kl,
predict_xstart,
rescale_timesteps,
rescale_learned_sigmas,
use_checkpoint,
use_scale_shift_norm,
resblock_updown,
use_fp16,
use_new_attention_order,
):
model = create_model(
image_size,
num_channels,
num_res_blocks,
channel_mult=channel_mult,
learn_sigma=learn_sigma,
class_cond=class_cond,
use_checkpoint=use_checkpoint,
attention_resolutions=attention_resolutions,
num_heads=num_heads,
num_head_channels=num_head_channels,
num_heads_upsample=num_heads_upsample,
use_scale_shift_norm=use_scale_shift_norm,
dropout=dropout,
resblock_updown=resblock_updown,
use_fp16=use_fp16,
use_new_attention_order=use_new_attention_order,
)
diffusion = create_gaussian_diffusion(
steps=diffusion_steps,
learn_sigma=learn_sigma,
noise_schedule=noise_schedule,
use_kl=use_kl,
predict_xstart=predict_xstart,
rescale_timesteps=rescale_timesteps,
rescale_learned_sigmas=rescale_learned_sigmas,
timestep_respacing=timestep_respacing,
)
return model, diffusion
def create_model(
image_size,
num_channels,
num_res_blocks,
channel_mult="",
learn_sigma=False,
class_cond=False,
use_checkpoint=False,
attention_resolutions="16",
num_heads=1,
num_head_channels=-1,
num_heads_upsample=-1,
use_scale_shift_norm=False,
dropout=0,
resblock_updown=False,
use_fp16=False,
use_new_attention_order=False,
):
if channel_mult == "":
if image_size == 512:
channel_mult = (0.5, 1, 1, 2, 2, 4, 4)
elif image_size == 256:
channel_mult = (1, 1, 2, 2, 4, 4)
elif image_size == 128:
channel_mult = (1, 1, 2, 3, 4)
elif image_size == 64:
channel_mult = (1, 2, 3, 4)
else:
raise ValueError(f"unsupported image size: {image_size}")
else:
channel_mult = tuple(int(ch_mult) for ch_mult in channel_mult.split(","))
attention_ds = []
for res in attention_resolutions.split(","):
attention_ds.append(image_size // int(res))
return UNetModel(
image_size=image_size,
in_channels=3,
model_channels=num_channels,
out_channels=(3 if not learn_sigma else 6),
num_res_blocks=num_res_blocks,
attention_resolutions=tuple(attention_ds),
dropout=dropout,
channel_mult=channel_mult,
num_classes=(NUM_CLASSES if class_cond else None),
use_checkpoint=use_checkpoint,
use_fp16=use_fp16,
num_heads=num_heads,
num_head_channels=num_head_channels,
num_heads_upsample=num_heads_upsample,
use_scale_shift_norm=use_scale_shift_norm,
resblock_updown=resblock_updown,
use_new_attention_order=use_new_attention_order,
)
def create_classifier_and_diffusion(
image_size,
classifier_use_fp16,
classifier_width,
classifier_depth,
classifier_attention_resolutions,
classifier_use_scale_shift_norm,
classifier_resblock_updown,
classifier_pool,
learn_sigma,
diffusion_steps,
noise_schedule,
timestep_respacing,
use_kl,
predict_xstart,
rescale_timesteps,
rescale_learned_sigmas,
):
classifier = create_classifier(
image_size,
classifier_use_fp16,
classifier_width,
classifier_depth,
classifier_attention_resolutions,
classifier_use_scale_shift_norm,
classifier_resblock_updown,
classifier_pool,
)
diffusion = create_gaussian_diffusion(
steps=diffusion_steps,
learn_sigma=learn_sigma,
noise_schedule=noise_schedule,
use_kl=use_kl,
predict_xstart=predict_xstart,
rescale_timesteps=rescale_timesteps,
rescale_learned_sigmas=rescale_learned_sigmas,
timestep_respacing=timestep_respacing,
)
return classifier, diffusion
def create_classifier(
image_size,
classifier_use_fp16,
classifier_width,
classifier_depth,
classifier_attention_resolutions,
classifier_use_scale_shift_norm,
classifier_resblock_updown,
classifier_pool,
):
if image_size == 512:
channel_mult = (0.5, 1, 1, 2, 2, 4, 4)
elif image_size == 256:
channel_mult = (1, 1, 2, 2, 4, 4)
elif image_size == 128:
channel_mult = (1, 1, 2, 3, 4)
elif image_size == 64:
channel_mult = (1, 2, 3, 4)
else:
raise ValueError(f"unsupported image size: {image_size}")
attention_ds = []
for res in classifier_attention_resolutions.split(","):
attention_ds.append(image_size // int(res))
return EncoderUNetModel(
image_size=image_size,
in_channels=3,
model_channels=classifier_width,
out_channels=1000,
num_res_blocks=classifier_depth,
attention_resolutions=tuple(attention_ds),
channel_mult=channel_mult,
use_fp16=classifier_use_fp16,
num_head_channels=64,
use_scale_shift_norm=classifier_use_scale_shift_norm,
resblock_updown=classifier_resblock_updown,
pool=classifier_pool,
)
def sr_model_and_diffusion_defaults():
res = model_and_diffusion_defaults()
res["large_size"] = 256
res["small_size"] = 64
arg_names = inspect.getfullargspec(sr_create_model_and_diffusion)[0]
for k in res.copy().keys():
if k not in arg_names:
del res[k]
return res
def sr_create_model_and_diffusion(
large_size,
small_size,
class_cond,
learn_sigma,
num_channels,
num_res_blocks,
num_heads,
num_head_channels,
num_heads_upsample,
attention_resolutions,
dropout,
diffusion_steps,
noise_schedule,
timestep_respacing,
use_kl,
predict_xstart,
rescale_timesteps,
rescale_learned_sigmas,
use_checkpoint,
use_scale_shift_norm,
resblock_updown,
use_fp16,
):
model = sr_create_model(
large_size,
small_size,
num_channels,
num_res_blocks,
learn_sigma=learn_sigma,
class_cond=class_cond,
use_checkpoint=use_checkpoint,
attention_resolutions=attention_resolutions,
num_heads=num_heads,
num_head_channels=num_head_channels,
num_heads_upsample=num_heads_upsample,
use_scale_shift_norm=use_scale_shift_norm,
dropout=dropout,
resblock_updown=resblock_updown,
use_fp16=use_fp16,
)
diffusion = create_gaussian_diffusion(
steps=diffusion_steps,
learn_sigma=learn_sigma,
noise_schedule=noise_schedule,
use_kl=use_kl,
predict_xstart=predict_xstart,
rescale_timesteps=rescale_timesteps,
rescale_learned_sigmas=rescale_learned_sigmas,
timestep_respacing=timestep_respacing,
)
return model, diffusion
def sr_create_model(
large_size,
small_size,
num_channels,
num_res_blocks,
learn_sigma,
class_cond,
use_checkpoint,
attention_resolutions,
num_heads,
num_head_channels,
num_heads_upsample,
use_scale_shift_norm,
dropout,
resblock_updown,
use_fp16,
):
_ = small_size # hack to prevent unused variable
if large_size == 512:
channel_mult = (1, 1, 2, 2, 4, 4)
elif large_size == 256:
channel_mult = (1, 1, 2, 2, 4, 4)
elif large_size == 64:
channel_mult = (1, 2, 3, 4)
else:
raise ValueError(f"unsupported large size: {large_size}")
attention_ds = []
for res in attention_resolutions.split(","):
attention_ds.append(large_size // int(res))
return SuperResModel(
image_size=large_size,
in_channels=3,
model_channels=num_channels,
out_channels=(3 if not learn_sigma else 6),
num_res_blocks=num_res_blocks,
attention_resolutions=tuple(attention_ds),
dropout=dropout,
channel_mult=channel_mult,
num_classes=(NUM_CLASSES if class_cond else None),
use_checkpoint=use_checkpoint,
num_heads=num_heads,
num_head_channels=num_head_channels,
num_heads_upsample=num_heads_upsample,
use_scale_shift_norm=use_scale_shift_norm,
resblock_updown=resblock_updown,
use_fp16=use_fp16,
)
def create_gaussian_diffusion(
*,
steps=1000,
learn_sigma=False,
sigma_small=False,
noise_schedule="linear",
use_kl=False,
predict_xstart=False,
rescale_timesteps=False,
rescale_learned_sigmas=False,
timestep_respacing="",
):
betas = gd.get_named_beta_schedule(noise_schedule, steps)
if use_kl:
loss_type = gd.LossType.RESCALED_KL
elif rescale_learned_sigmas:
loss_type = gd.LossType.RESCALED_MSE
else:
loss_type = gd.LossType.MSE
if not timestep_respacing:
timestep_respacing = [steps]
return SpacedDiffusion(
use_timesteps=space_timesteps(steps, timestep_respacing),
betas=betas,
model_mean_type=(
gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X
),
model_var_type=(
(
gd.ModelVarType.FIXED_LARGE
if not sigma_small
else gd.ModelVarType.FIXED_SMALL
)
if not learn_sigma
else gd.ModelVarType.LEARNED_RANGE
),
loss_type=loss_type,
rescale_timesteps=rescale_timesteps,
)
def add_dict_to_argparser(parser, default_dict):
for k, v in default_dict.items():
v_type = type(v)
if v is None:
v_type = str
elif isinstance(v, bool):
v_type = str2bool
parser.add_argument(f"--{k}", default=v, type=v_type)
def args_to_dict(args, keys):
return {k: getattr(args, k) for k in keys}
def str2bool(v):
"""
https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse
"""
if isinstance(v, bool):
return v
if v.lower() in ("yes", "true", "t", "y", "1"):
return True
elif v.lower() in ("no", "false", "f", "n", "0"):
return False
else:
raise argparse.ArgumentTypeError("boolean value expected")
================================================
FILE: guided_diffusion/train_util.py
================================================
import copy
import functools
import os
import blobfile as bf
import torch as th
import torch.distributed as dist
from torch.nn.parallel.distributed import DistributedDataParallel as DDP
from torch.optim import AdamW
from . import dist_util, logger
from .fp16_util import MixedPrecisionTrainer
from .nn import update_ema
from .resample import LossAwareSampler, UniformSampler
# For ImageNet experiments, this was a good default value.
# We found that the lg_loss_scale quickly climbed to
# 20-21 within the first ~1K steps of training.
INITIAL_LOG_LOSS_SCALE = 20.0
class TrainLoop:
def __init__(
self,
*,
model,
diffusion,
data,
batch_size,
microbatch,
lr,
ema_rate,
log_interval,
save_interval,
resume_checkpoint,
use_fp16=False,
fp16_scale_growth=1e-3,
schedule_sampler=None,
weight_decay=0.0,
lr_anneal_steps=0,
):
self.model = model
self.diffusion = diffusion
self.data = data
self.batch_size = batch_size
self.microbatch = microbatch if microbatch > 0 else batch_size
self.lr = lr
self.ema_rate = (
[ema_rate]
if isinstance(ema_rate, float)
else [float(x) for x in ema_rate.split(",")]
)
self.log_interval = log_interval
self.save_interval = save_interval
self.resume_checkpoint = resume_checkpoint
self.use_fp16 = use_fp16
self.fp16_scale_growth = fp16_scale_growth
self.schedule_sampler = schedule_sampler or UniformSampler(diffusion)
self.weight_decay = weight_decay
self.lr_anneal_steps = lr_anneal_steps
self.step = 0
self.resume_step = 0
self.global_batch = self.batch_size * dist.get_world_size()
self.sync_cuda = th.cuda.is_available()
self._load_and_sync_parameters()
self.mp_trainer = MixedPrecisionTrainer(
model=self.model,
use_fp16=self.use_fp16,
fp16_scale_growth=fp16_scale_growth,
)
self.opt = AdamW(
self.mp_trainer.master_params, lr=self.lr, weight_decay=self.weight_decay
)
if self.resume_step:
self._load_optimizer_state()
# Model was resumed, either due to a restart or a checkpoint
# being specified at the command line.
self.ema_params = [
self._load_ema_parameters(rate) for rate in self.ema_rate
]
else:
self.ema_params = [
copy.deepcopy(self.mp_trainer.master_params)
for _ in range(len(self.ema_rate))
]
if th.cuda.is_available():
self.use_ddp = True
self.ddp_model = DDP(
self.model,
device_ids=[dist_util.dev()],
output_device=dist_util.dev(),
broadcast_buffers=False,
bucket_cap_mb=128,
find_unused_parameters=False,
)
else:
if dist.get_world_size() > 1:
logger.warn(
"Distributed training requires CUDA. "
"Gradients will not be synchronized properly!"
)
self.use_ddp = False
self.ddp_model = self.model
def _load_and_sync_parameters(self):
resume_checkpoint = find_resume_checkpoint() or self.resume_checkpoint
if resume_checkpoint:
self.resume_step = parse_resume_step_from_filename(resume_checkpoint)
if dist.get_rank() == 0:
logger.log(f"loading model from checkpoint: {resume_checkpoint}...")
self.model.load_state_dict(
dist_util.load_state_dict(
resume_checkpoint, map_location=dist_util.dev()
)
)
dist_util.sync_params(self.model.parameters())
def _load_ema_parameters(self, rate):
ema_params = copy.deepcopy(self.mp_trainer.master_params)
main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint
ema_checkpoint = find_ema_checkpoint(main_checkpoint, self.resume_step, rate)
if ema_checkpoint:
if dist.get_rank() == 0:
logger.log(f"loading EMA from checkpoint: {ema_checkpoint}...")
state_dict = dist_util.load_state_dict(
ema_checkpoint, map_location=dist_util.dev()
)
ema_params = self.mp_trainer.state_dict_to_master_params(state_dict)
dist_util.sync_params(ema_params)
return ema_params
def _load_optimizer_state(self):
main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint
opt_checkpoint = bf.join(
bf.dirname(main_checkpoint), f"opt{self.resume_step:06}.pt"
)
if bf.exists(opt_checkpoint):
logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}")
state_dict = dist_util.load_state_dict(
opt_checkpoint, map_location=dist_util.dev()
)
self.opt.load_state_dict(state_dict)
def run_loop(self):
while (
not self.lr_anneal_steps
or self.step + self.resume_step < self.lr_anneal_steps
):
batch, cond = next(self.data)
self.run_step(batch, cond)
if self.step % self.log_interval == 0:
logger.dumpkvs()
if self.step % self.save_interval == 0:
self.save()
# Run for a finite amount of time in integration tests.
if os.environ.get("DIFFUSION_TRAINING_TEST", "") and self.step > 0:
return
self.step += 1
# Save the last checkpoint if it wasn't already saved.
if (self.step - 1) % self.save_interval != 0:
self.save()
def run_step(self, batch, cond):
self.forward_backward(batch, cond)
took_step = self.mp_trainer.optimize(self.opt)
if took_step:
self._update_ema()
self._anneal_lr()
self.log_step()
def forward_backward(self, batch, cond):
self.mp_trainer.zero_grad()
for i in range(0, batch.shape[0], self.microbatch):
micro = batch[i : i + self.microbatch].to(dist_util.dev())
micro_cond = {
k: v[i : i + self.microbatch].to(dist_util.dev())
for k, v in cond.items()
}
last_batch = (i + self.microbatch) >= batch.shape[0]
t, weights = self.schedule_sampler.sample(micro.shape[0], dist_util.dev())
# print("TrainLoop/forward_backward() t:", t)
assert t[0] >= 0, f"train_util.py t[0] : {t[0]}"
compute_losses = functools.partial(
self.diffusion.training_losses,
self.ddp_model,
micro,
t,
model_kwargs=micro_cond,
)
if last_batch or not self.use_ddp:
losses = compute_losses()
else:
with self.ddp_model.no_sync():
losses = compute_losses()
if isinstance(self.schedule_sampler, LossAwareSampler):
self.schedule_sampler.update_with_local_losses(
t, losses["loss"].detach()
)
loss = (losses["loss"] * weights).mean()
log_loss_dict(
self.diffusion, t, {k: v * weights for k, v in losses.items()}
)
self.mp_trainer.backward(loss)
def _update_ema(self):
for rate, params in zip(self.ema_rate, self.ema_params):
update_ema(params, self.mp_trainer.master_params, rate=rate)
def _anneal_lr(self):
if not self.lr_anneal_steps:
return
frac_done = (self.step + self.resume_step) / self.lr_anneal_steps
lr = self.lr * (1 - frac_done)
for param_group in self.opt.param_groups:
param_group["lr"] = lr
def log_step(self):
logger.logkv("step", self.step + self.resume_step)
logger.logkv("samples", (self.step + self.resume_step + 1) * self.global_batch)
def save(self):
def save_checkpoint(rate, params):
state_dict = self.mp_trainer.master_params_to_state_dict(params)
if dist.get_rank() == 0:
logger.log(f"saving model {rate}...")
if not rate:
filename = f"model{(self.step+self.resume_step):06d}.pt"
else:
filename = f"ema_{rate}_{(self.step+self.resume_step):06d}.pt"
with bf.BlobFile(bf.join(get_blob_logdir(), filename), "wb") as f:
th.save(state_dict, f)
save_checkpoint(0, self.mp_trainer.master_params)
for rate, params in zip(self.ema_rate, self.ema_params):
save_checkpoint(rate, params)
if dist.get_rank() == 0:
with bf.BlobFile(
bf.join(get_blob_logdir(), f"opt{(self.step+self.resume_step):06d}.pt"),
"wb",
) as f:
th.save(self.opt.state_dict(), f)
dist.barrier()
def parse_resume_step_from_filename(filename):
"""
Parse filenames of the form path/to/modelNNNNNN.pt, where NNNNNN is the
checkpoint's number of steps.
"""
split = filename.split("model")
if len(split) < 2:
return 0
split1 = split[-1].split(".")[0]
try:
return int(split1)
except ValueError:
return 0
def get_blob_logdir():
# You can change this to be a separate path to save checkpoints to
# a blobstore or some external drive.
return logger.get_dir()
def find_resume_checkpoint():
# On your infrastructure, you may want to override this to automatically
# discover the latest checkpoint on your blob storage, etc.
return None
def find_ema_checkpoint(main_checkpoint, step, rate):
if main_checkpoint is None:
return None
filename = f"ema_{rate}_{(step):06d}.pt"
path = bf.join(bf.dirname(main_checkpoint), filename)
if bf.exists(path):
return path
return None
def log_loss_dict(diffusion, ts, losses):
for key, values in losses.items():
logger.logkv_mean(key, values.mean().item())
# Log the quantiles (four quartiles, in particular).
for sub_t, sub_loss in zip(ts.cpu().numpy(), values.detach().cpu().numpy()):
quartile = int(4 * sub_t / diffusion.num_timesteps)
logger.logkv_mean(f"{key}_q{quartile}", sub_loss)
================================================
FILE: guided_diffusion/unet.py
================================================
from abc import abstractmethod
import math
import numpy as np
import torch as th
import torch.nn as nn
import torch.nn.functional as F
from .fp16_util import convert_module_to_f16, convert_module_to_f32
from .nn import (
checkpoint,
conv_nd,
linear,
avg_pool_nd,
zero_module,
normalization,
timestep_embedding,
)
class AttentionPool2d(nn.Module):
"""
Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py
"""
def __init__(
self,
spacial_dim: int,
embed_dim: int,
num_heads_channels: int,
output_dim: int = None,
):
super().__init__()
self.positional_embedding = nn.Parameter(
th.randn(embed_dim, spacial_dim ** 2 + 1) / embed_dim ** 0.5
)
self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1)
self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1)
self.num_heads = embed_dim // num_heads_channels
self.attention = QKVAttention(self.num_heads)
def forward(self, x):
b, c, *_spatial = x.shape
x = x.reshape(b, c, -1) # NC(HW)
x = th.cat([x.mean(dim=-1, keepdim=True), x], dim=-1) # NC(HW+1)
x = x + self.positional_embedding[None, :, :].to(x.dtype) # NC(HW+1)
x = self.qkv_proj(x)
x = self.attention(x)
x = self.c_proj(x)
return x[:, :, 0]
class TimestepBlock(nn.Module):
"""
Any module where forward() takes timestep embeddings as a second argument.
"""
@abstractmethod
def forward(self, x, emb):
"""
Apply the module to `x` given `emb` timestep embeddings.
"""
class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
"""
A sequential module that passes timestep embeddings to the children that
support it as an extra input.
"""
def forward(self, x, emb):
for layer in self:
if isinstance(layer, TimestepBlock):
x = layer(x, emb)
else:
x = layer(x)
return x
class Upsample(nn.Module):
"""
An upsampling layer with an optional convolution.
:param channels: channels in the inputs and outputs.
:param use_conv: a bool determining if a convolution is applied.
:param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
upsampling occurs in the inner-two dimensions.
"""
def __init__(self, channels, use_conv, dims=2, out_channels=None):
super().__init__()
self.channels = channels
self.out_channels = out_channels or channels
self.use_conv = use_conv
self.dims = dims
if use_conv:
self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=1)
def forward(self, x):
assert x.shape[1] == self.channels
if self.dims == 3:
x = F.interpolate(
x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
)
else:
x = F.interpolate(x, scale_factor=2, mode="nearest")
if self.use_conv:
x = self.conv(x)
return x
class Downsample(nn.Module):
"""
A downsampling layer with an optional convolution.
:param channels: channels in the inputs and outputs.
:param use_conv: a bool determining if a convolution is applied.
:param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
downsampling occurs in the inner-two dimensions.
"""
def __init__(self, channels, use_conv, dims=2, out_channels=None):
super().__init__()
self.channels = channels
self.out_channels = out_channels or channels
self.use_conv = use_conv
self.dims = dims
stride = 2 if dims != 3 else (1, 2, 2)
if use_conv:
self.op = conv_nd(
dims, self.channels, self.out_channels, 3, stride=stride, padding=1
)
else:
assert self.channels == self.out_channels
self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
def forward(self, x):
assert x.shape[1] == self.channels
return self.op(x)
class ResBlock(TimestepBlock):
"""
A residual block that can optionally change the number of channels.
:param channels: the number of input channels.
:param emb_channels: the number of timestep embedding channels.
:param dropout: the rate of dropout.
:param out_channels: if specified, the number of out channels.
:param use_conv: if True and out_channels is specified, use a spatial
convolution instead of a smaller 1x1 convolution to change the
channels in the skip connection.
:param dims: determines if the signal is 1D, 2D, or 3D.
:param use_checkpoint: if True, use gradient checkpointing on this module.
:param up: if True, use this block for upsampling.
:param down: if True, use this block for downsampling.
"""
def __init__(
self,
channels,
emb_channels,
dropout,
out_channels=None,
use_conv=False,
use_scale_shift_norm=False,
dims=2,
use_checkpoint=False,
up=False,
down=False,
):
super().__init__()
self.channels = channels
self.emb_channels = emb_channels
self.dropout = dropout
self.out_channels = out_channels or channels
self.use_conv = use_conv
self.use_checkpoint = use_checkpoint
self.use_scale_shift_norm = use_scale_shift_norm
self.in_layers = nn.Sequential(
normalization(channels),
nn.SiLU(),
conv_nd(dims, channels, self.out_channels, 3, padding=1),
)
self.updown = up or down
if up:
self.h_upd = Upsample(channels, False, dims)
self.x_upd = Upsample(channels, False, dims)
elif down:
self.h_upd = Downsample(channels, False, dims)
self.x_upd = Downsample(channels, False, dims)
else:
self.h_upd = self.x_upd = nn.Identity()
self.emb_layers = nn.Sequential(
nn.SiLU(),
linear(
emb_channels,
2 * self.out_channels if use_scale_shift_norm else self.out_channels,
),
)
self.out_layers = nn.Sequential(
normalization(self.out_channels),
nn.SiLU(),
nn.Dropout(p=dropout),
zero_module(
conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)
),
)
if self.out_channels == channels:
self.skip_connection = nn.Identity()
elif use_conv:
self.skip_connection = conv_nd(
dims, channels, self.out_channels, 3, padding=1
)
else:
self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
def forward(self, x, emb):
"""
Apply the block to a Tensor, conditioned on a timestep embedding.
:param x: an [N x C x ...] Tensor of features.
:param emb: an [N x emb_channels] Tensor of timestep embeddings.
:return: an [N x C x ...] Tensor of outputs.
"""
return checkpoint(
self._forward, (x, emb), self.parameters(), self.use_checkpoint
)
def _forward(self, x, emb):
if self.updown:
in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
h = in_rest(x)
h = self.h_upd(h)
x = self.x_upd(x)
h = in_conv(h)
else:
h = self.in_layers(x)
emb_out = self.emb_layers(emb).type(h.dtype)
while len(emb_out.shape) < len(h.shape):
emb_out = emb_out[..., None]
if self.use_scale_shift_norm:
out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
scale, shift = th.chunk(emb_out, 2, dim=1)
h = out_norm(h) * (1 + scale) + shift
h = out_rest(h)
else:
h = h + emb_out
h = self.out_layers(h)
return self.skip_connection(x) + h
class AttentionBlock(nn.Module):
"""
An attention block that allows spatial positions to attend to each other.
Originally ported from here, but adapted to the N-d case.
https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
"""
def __init__(
self,
channels,
num_heads=1,
num_head_channels=-1,
use_checkpoint=False,
use_new_attention_order=False,
):
super().__init__()
self.channels = channels
if num_head_channels == -1:
self.num_heads = num_heads
else:
assert (
channels % num_head_channels == 0
), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
self.num_heads = channels // num_head_channels
self.use_checkpoint = use_checkpoint
self.norm = normalization(channels)
self.qkv = conv_nd(1, channels, channels * 3, 1)
if use_new_attention_order:
# split qkv before split heads
self.attention = QKVAttention(self.num_heads)
else:
# split heads before split qkv
self.attention = QKVAttentionLegacy(self.num_heads)
self.proj_out = zero_module(conv_nd(1, channels, channels, 1))
def forward(self, x):
return checkpoint(self._forward, (x,), self.parameters(), True)
def _forward(self, x):
b, c, *spatial = x.shape
x = x.reshape(b, c, -1)
qkv = self.qkv(self.norm(x))
h = self.attention(qkv)
h = self.proj_out(h)
return (x + h).reshape(b, c, *spatial)
def count_flops_attn(model, _x, y):
"""
A counter for the `thop` package to count the operations in an
attention operation.
Meant to be used like:
macs, params = thop.profile(
model,
inputs=(inputs, timestamps),
custom_ops={QKVAttention: QKVAttention.count_flops},
)
"""
b, c, *spatial = y[0].shape
num_spatial = int(np.prod(spatial))
# We perform two matmuls with the same number of ops.
# The first computes the weight matrix, the second computes
# the combination of the value vectors.
matmul_ops = 2 * b * (num_spatial ** 2) * c
model.total_ops += th.DoubleTensor([matmul_ops])
class QKVAttentionLegacy(nn.Module):
"""
A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping
"""
def __init__(self, n_heads):
super().__init__()
self.n_heads = n_heads
def forward(self, qkv):
"""
Apply QKV attention.
:param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
:return: an [N x (H * C) x T] tensor after attention.
"""
bs, width, length = qkv.shape
assert width % (3 * self.n_heads) == 0
ch = width // (3 * self.n_heads)
q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
scale = 1 / math.sqrt(math.sqrt(ch))
weight = th.einsum(
"bct,bcs->bts", q * scale, k * scale
) # More stable with f16 than dividing afterwards
weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
a = th.einsum("bts,bcs->bct", weight, v)
return a.reshape(bs, -1, length)
@staticmethod
def count_flops(model, _x, y):
return count_flops_attn(model, _x, y)
class QKVAttention(nn.Module):
"""
A module which performs QKV attention and splits in a different order.
"""
def __init__(self, n_heads):
super().__init__()
self.n_heads = n_heads
def forward(self, qkv):
"""
Apply QKV attention.
:param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs.
:return: an [N x (H * C) x T] tensor after attention.
"""
bs, width, length = qkv.shape
assert width % (3 * self.n_heads) == 0
ch = width // (3 * self.n_heads)
q, k, v = qkv.chunk(3, dim=1)
scale = 1 / math.sqrt(math.sqrt(ch))
weight = th.einsum(
"bct,bcs->bts",
(q * scale).view(bs * self.n_heads, ch, length),
(k * scale).view(bs * self.n_heads, ch, length),
) # More stable with f16 than dividing afterwards
weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
a = th.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length))
return a.reshape(bs, -1, length)
@staticmethod
def count_flops(model, _x, y):
return count_flops_attn(model, _x, y)
class UNetModel(nn.Module):
"""
The full UNet model with attention and timestep embedding.
:param in_channels: channels in the input Tensor.
:param model_channels: base channel count for the model.
:param out_channels: channels in the output Tensor.
:param num_res_blocks: number of residual blocks per downsample.
:param attention_resolutions: a collection of downsample rates at which
attention will take place. May be a set, list, or tuple.
For example, if this contains 4, then at 4x downsampling, attention
will be used.
:param dropout: the dropout probability.
:param channel_mult: channel multiplier for each level of the UNet.
:param conv_resample: if True, use learned convolutions for upsampling and
downsampling.
:param dims: determines if the signal is 1D, 2D, or 3D.
:param num_classes: if specified (as an int), then this model will be
class-conditional with `num_classes` classes.
:param use_checkpoint: use gradient checkpointing to reduce memory usage.
:param num_heads: the number of attention heads in each attention layer.
:param num_heads_channels: if specified, ignore num_heads and instead use
a fixed channel width per attention head.
:param num_heads_upsample: works with num_heads to set a different number
of heads for upsampling. Deprecated.
:param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
:param resblock_updown: use residual blocks for up/downsampling.
:param use_new_attention_order: use a different attention pattern for potentially
increased efficiency.
"""
def __init__(
self,
image_size,
in_channels,
model_channels,
out_channels,
blur,
num_res_blocks = 3,
attention_resolutions = [4,8],
dropout=0,
channel_mult=(1, 2, 3, 4),
conv_resample=True,
dims=2,
num_classes=None,
use_checkpoint=False,
use_fp16=False,
num_heads=4,
num_head_channels=-1,
num_heads_upsample=-1,
use_scale_shift_norm=True,
resblock_updown=False,
use_new_attention_order=False,
freq_feat = False
):
super().__init__()
if num_heads_upsample == -1:
num_heads_upsample = num_heads
self.image_size = image_size
self.in_channels = in_channels
self.model_channels = model_channels
self.out_channels = out_channels
self.num_res_blocks = num_res_blocks
self.attention_resolutions = attention_resolutions
self.dropout = dropout
self.channel_mult = channel_mult
self.conv_resample = conv_resample
self.num_classes = num_classes
self.use_checkpoint = use_checkpoint
self.dtype = th.float16 if use_fp16 else th.float32
self.num_heads = num_heads
self.num_head_channels = num_head_channels
self.num_heads_upsample = num_heads_upsample
self.freq_feat = freq_feat
self.blur = blur
time_embed_dim = model_channels * 4
self.time_embed = nn.Sequential(
linear(model_channels, time_embed_dim),
nn.SiLU(),
linear(time_embed_dim, time_embed_dim),
)
if freq_feat:
self.in_channels += 3
in_channels += 3
if self.num_classes is not None:
self.label_emb = nn.Embedding(num_classes, time_embed_dim)
ch = input_ch = int(channel_mult[0] * model_channels)
self.input_blocks = nn.ModuleList(
[TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))]
)
print("input_blocks", self.input_blocks[0][0].weight.shape)
if freq_feat:
# Set the weights to the frequency features to zeros
self.input_blocks[0][0].weight.data[:, 3: :, :] = 0
# for j in range(6):
# w = self.input_blocks[0][0].weight[:, j, :, :]
# print(f"j: {j} w.mean(): {w.mean()} w.std(): {w.std()}")
# raise NotImplementedError
self._feature_size = ch
input_block_chans = [ch]
ds = 1
for level, mult in enumerate(channel_mult):
for _ in range(num_res_blocks):
layers = [
ResBlock(
ch,
time_embed_dim,
dropout,
out_channels=int(mult * model_channels),
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
)
]
ch = int(mult * model_channels)
if ds in attention_resolutions:
layers.append(
AttentionBlock(
ch,
use_checkpoint=use_checkpoint,
num_heads=num_heads,
num_head_channels=num_head_channels,
use_new_attention_order=use_new_attention_order,
)
)
self.input_blocks.append(TimestepEmbedSequential(*layers))
self._feature_size += ch
input_block_chans.append(ch)
if level != len(channel_mult) - 1:
out_ch = ch
self.input_blocks.append(
TimestepEmbedSequential(
ResBlock(
ch,
time_embed_dim,
dropout,
out_channels=out_ch,
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
down=True,
)
if resblock_updown
else Downsample(
ch, conv_resample, dims=dims, out_channels=out_ch
)
)
)
ch = out_ch
input_block_chans.append(ch)
ds *= 2
self._feature_size += ch
self.middle_block = TimestepEmbedSequential(
ResBlock(
ch,
time_embed_dim,
dropout,
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
),
AttentionBlock(
ch,
use_checkpoint=use_checkpoint,
num_heads=num_heads,
num_head_channels=num_head_channels,
use_new_attention_order=use_new_attention_order,
),
ResBlock(
ch,
time_embed_dim,
dropout,
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
),
)
self._feature_size += ch
self.output_blocks = nn.ModuleList([])
for level, mult in list(enumerate(channel_mult))[::-1]:
for i in range(num_res_blocks + 1):
ich = input_block_chans.pop()
layers = [
ResBlock(
ch + ich,
time_embed_dim,
dropout,
out_channels=int(model_channels * mult),
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
)
]
ch = int(model_channels * mult)
if ds in attention_resolutions:
layers.append(
AttentionBlock(
ch,
use_checkpoint=use_checkpoint,
num_heads=num_heads_upsample,
num_head_channels=num_head_channels,
use_new_attention_order=use_new_attention_order,
)
)
if level and i == num_res_blocks:
out_ch = ch
layers.append(
ResBlock(
ch,
time_embed_dim,
dropout,
out_channels=out_ch,
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
up=True,
)
if resblock_updown
else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)
)
ds //= 2
self.output_blocks.append(TimestepEmbedSequential(*layers))
self._feature_size += ch
self.out = nn.Sequential(
normalization(ch),
nn.SiLU(),
zero_module(conv_nd(dims, input_ch, out_channels, 3, padding=1)),
)
def convert_to_fp16(self):
"""
Convert the torso of the model to float16.
"""
self.input_blocks.apply(convert_module_to_f16)
self.middle_block.apply(convert_module_to_f16)
self.output_blocks.apply(convert_module_to_f16)
def convert_to_fp32(self):
"""
Convert the torso of the model to float32.
"""
self.input_blocks.apply(convert_module_to_f32)
self.middle_block.apply(convert_module_to_f32)
self.output_blocks.apply(convert_module_to_f32)
def forward(self, x, timesteps, y=None):
"""
Apply the model to an input batch.
:param x: an [N x C x ...] Tensor of inputs.
:param timesteps: a 1-D batch of timesteps.
:param y: an [N] Tensor of labels, if class-conditional.
:return: an [N x C x ...] Tensor of outputs.
"""
assert (y is not None) == (
self.num_classes is not None
), "must specify y if and only if the model is class-conditional"
# timesteps = [timesteps] * x.shape[0]
# timesteps = th.tensor(timesteps).to(x.device)
if self.freq_feat:
x_freq = self.blur.Ut(x).view(x.shape)
x = th.cat([x, x_freq], dim=1)
hs = []
emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
if self.num_classes is not None:
assert y.shape == (x.shape[0],)
emb = emb + self.label_emb(y)
h = x.type(self.dtype)
for module in self.input_blocks:
h = module(h, emb)
hs.append(h)
h = self.middle_block(h, emb)
for module in self.output_blocks:
h = th.cat([h, hs.pop()], dim=1)
h = module(h, emb)
h = h.type(x.dtype)
return self.out(h)
class SuperResModel(UNetModel):
"""
A UNetModel that performs super-resolution.
Expects an extra kwarg `low_res` to condition on a low-resolution image.
"""
def __init__(self, image_size, in_channels, *args, **kwargs):
super().__init__(image_size, in_channels * 2, *args, **kwargs)
def forward(self, x, timesteps, low_res=None, **kwargs):
_, _, new_height, new_width = x.shape
upsampled = F.interpolate(low_res, (new_height, new_width), mode="bilinear")
x = th.cat([x, upsampled], dim=1)
return super().forward(x, timesteps, **kwargs)
class EncoderUNetModel(nn.Module):
"""
The half UNet model with attention and timestep embedding.
For usage, see UNet.
"""
def __init__(
self,
image_size,
in_channels,
model_channels,
out_channels,
num_res_blocks,
attention_resolutions,
dropout=0,
channel_mult=(1, 2, 4, 8),
conv_resample=True,
dims=2,
use_checkpoint=False,
use_fp16=False,
num_heads=1,
num_head_channels=-1,
num_heads_upsample=-1,
use_scale_shift_norm=False,
resblock_updown=False,
use_new_attention_order=False,
pool="adaptive",
):
super().__init__()
if num_heads_upsample == -1:
num_heads_upsample = num_heads
self.in_channels = in_channels
self.model_channels = model_channels
self.out_channels = out_channels
self.num_res_blocks = num_res_blocks
self.attention_resolutions = attention_resolutions
self.dropout = dropout
self.channel_mult = channel_mult
self.conv_resample = conv_resample
self.use_checkpoint = use_checkpoint
self.dtype = th.float16 if use_fp16 else th.float32
self.num_heads = num_heads
self.num_head_channels = num_head_channels
self.num_heads_upsample = num_heads_upsample
time_embed_dim = model_channels * 4
self.time_embed = nn.Sequential(
linear(model_channels, time_embed_dim),
nn.SiLU(),
linear(time_embed_dim, time_embed_dim),
)
ch = int(channel_mult[0] * model_channels)
self.input_blocks = nn.ModuleList(
[TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))]
)
self._feature_size = ch
input_block_chans = [ch]
ds = 1
for level, mult in enumerate(channel_mult):
for _ in range(num_res_blocks):
layers = [
ResBlock(
ch,
time_embed_dim,
dropout,
out_channels=int(mult * model_channels),
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
)
]
ch = int(mult * model_channels)
if ds in attention_resolutions:
layers.append(
AttentionBlock(
ch,
use_checkpoint=use_checkpoint,
num_heads=num_heads,
num_head_channels=num_head_channels,
use_new_attention_order=use_new_attention_order,
)
)
self.input_blocks.append(TimestepEmbedSequential(*layers))
self._feature_size += ch
input_block_chans.append(ch)
if level != len(channel_mult) - 1:
out_ch = ch
self.input_blocks.append(
TimestepEmbedSequential(
ResBlock(
ch,
time_embed_dim,
dropout,
out_channels=out_ch,
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
down=True,
)
if resblock_updown
else Downsample(
ch, conv_resample, dims=dims, out_channels=out_ch
)
)
)
ch = out_ch
input_block_chans.append(ch)
ds *= 2
self._feature_size += ch
self.middle_block = TimestepEmbedSequential(
ResBlock(
ch,
time_embed_dim,
dropout,
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
),
AttentionBlock(
ch,
use_checkpoint=use_checkpoint,
num_heads=num_heads,
num_head_channels=num_head_channels,
use_new_attention_order=use_new_attention_order,
),
ResBlock(
ch,
time_embed_dim,
dropout,
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
),
)
self._feature_size += ch
self.pool = pool
if pool == "adaptive":
self.out = nn.Sequential(
normalization(ch),
nn.SiLU(),
nn.AdaptiveAvgPool2d((1, 1)),
zero_module(conv_nd(dims, ch, out_channels, 1)),
nn.Flatten(),
)
elif pool == "attention":
assert num_head_channels != -1
self.out = nn.Sequential(
normalization(ch),
nn.SiLU(),
AttentionPool2d(
(image_size // ds), ch, num_head_channels, out_channels
),
)
elif pool == "spatial":
self.out = nn.Sequential(
nn.Linear(self._feature_size, 2048),
nn.ReLU(),
nn.Linear(2048, self.out_channels),
)
elif pool == "spatial_v2":
self.out = nn.Sequential(
nn.Linear(self._feature_size, 2048),
normalization(2048),
nn.SiLU(),
nn.Linear(2048, self.out_channels),
)
else:
raise NotImplementedError(f"Unexpected {pool} pooling")
def convert_to_fp16(self):
"""
Convert the torso of the model to float16.
"""
self.input_blocks.apply(convert_module_to_f16)
self.middle_block.apply(convert_module_to_f16)
def convert_to_fp32(self):
"""
Convert the torso of the model to float32.
"""
self.input_blocks.apply(convert_module_to_f32)
self.middle_block.apply(convert_module_to_f32)
def forward(self, x, timesteps):
"""
Apply the model to an input batch.
:param x: an [N x C x ...] Tensor of inputs.
:param timesteps: a 1-D batch of timesteps.
:return: an [N x K] Tensor of outputs.
"""
emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
results = []
h = x.type(self.dtype)
for module in self.input_blocks:
h = module(h, emb)
if self.pool.startswith("spatial"):
results.append(h.type(x.dtype).mean(dim=(2, 3)))
h = self.middle_block(h, emb)
if self.pool.startswith("spatial"):
results.append(h.type(x.dtype).mean(dim=(2, 3)))
h = th.cat(results, axis=-1)
return self.out(h)
else:
h = h.type(x.dtype)
return self.out(h)
================================================
FILE: main.py
================================================
import torch
from PIL import Image
from collections import defaultdict
import torch.nn.functional as TF
import torchvision.datasets as dsets
from torchvision import transforms
import numpy as np
import torch.optim as optim
import torchvision
import matplotlib.pyplot as plt
import utils
from guided_diffusion.unet import UNetModel
import math
from tensorboardX import SummaryWriter
import os
import json
from collections import namedtuple
import argparse
from torchvision.utils import save_image
from tqdm import tqdm
from blur_diffusion import Deblurring, ForwardBlurIncreasing, gaussian_kernel_1d
from utils import normalize_np, clear
from EMA import EMA
from torch.nn import DataParallel
from fid import FID
from scipy.integrate import solve_ivp
parser = argparse.ArgumentParser(description='Configs')
parser.add_argument('--gpu', type=str, help='gpu num')
parser.add_argument('--dataset', type=str, help='cifar10 / mnist')
parser.add_argument('--name', type=str, help='Saving directory name')
parser.add_argument('--ckpt', default='', type=str, help='UNet checkpoint')
parser.add_argument('--bsize', default=16, type=int, help='batchsize')
parser.add_argument('--N', default=1000, type=int, help='Max diffusion timesteps')
parser.add_argument('--sig', default=0.4, type=float, help='sigma value for blur kernel')
parser.add_argument('--sig_min', default=0, type=float, help='sigma value for blur kernel')
parser.add_argument('--sig_max', default=0.1, type=float, help='sigma value for blur kernel')
parser.add_argument('--lr', default=0.00005, type=float, help='learning rate')
parser.add_argument('--noise_schedule', default='linear', type=str, help='Type of noise schedule to use')
parser.add_argument('--betamin', default=0.0001, type=float, help='beta (min). get_score(1) can diverge if this is too low.')
parser.add_argument('--betamax', default=0.02, type=float, help='beta (max)')
parser.add_argument('--fromprior', default=True, type=bool, help='start sampling from prior')
parser.add_argument('--gtscore', action='store_true', help='Use ground truth score for reverse diffusion')
parser.add_argument('--max_iter', default=1500000, type=int, help='max iterations')
parser.add_argument('--eval_iter', default=10000, type=int, help='eval iterations')
parser.add_argument('--fid_iter', default=50000, type=int, help='eval iterations')
parser.add_argument('--fid_num_samples', default=10000, type=int, help='eval iterations')
parser.add_argument('--fid_bsize', default=32, type=int, help='eval iterations')
parser.add_argument('--loss_type', type=str, default = 'eps_simple', choices=['sm_simple', 'eps_simple', 'sm_exact', 'std_matching'])
parser.add_argument('--f_type', type=str, default = 'linear', choices=['linear', 'log', 'quadratic', 'cubic', 'quartic', 'triangular'])
parser.add_argument('--dropout', default=0, type=float, help='dropout')
# EMA, save
parser.add_argument('--use_ema', action='store_true',
help='use EMA or not')
parser.add_argument('--inference', action='store_true')
parser.add_argument('--freq_feat', action='store_true', help = "concat Utx_i")
parser.add_argument('--ode', action='store_true', help = "ODE fast sampler")
parser.add_argument('--ema_decay', type=float, default=0.9999, help='decay rate for EMA')
parser.add_argument('--save_every', type=int, default=50000, help='How often we wish to save ckpts')
opt = parser.parse_args()
dataset = opt.dataset
device = torch.device(f'cuda:{opt.gpu}')
device = torch.device('cuda')
print("N:", opt.N)
N = opt.N
bsize = opt.bsize
beta_min = opt.betamin
beta_max = opt.betamax
sig = opt.sig
if opt.gtscore:
opt.fromprior = True
if dataset == 'mnist':
train_dir = '/home/ubuntu/code/sangyoon/datasets/mnist_train'
dataset_train = dsets.MNIST(root='/home/ubuntu/code/sangyoon/datasets/mnist_train',
train=True,
transform=transforms.Compose([transforms.Resize((32, 32)), transforms.ToTensor()]),
download=True)
dataset_test = dsets.MNIST(root='/home/ubuntu/code/sangyoon/datasets/mnist_test',
train=False,
transform=transforms.Compose([transforms.Resize((32, 32)), transforms.ToTensor()]),
download=True)
elif dataset == 'cifar10':
train_dir = "/home/ubuntu/code/sangyoon/datasets/cifar10_train/png"
dataset_train = dsets.CIFAR10(root='../datasets/cifar10_train',
train=True,
transform=transforms.Compose([transforms.RandomHorizontalFlip(0.5), transforms.ToTensor()]),
download=True
)
dataset_test = dsets.CIFAR10(root='../datasets/cifar10_test',
train=False,
transform=transforms.Compose([transforms.RandomHorizontalFlip(0.5), transforms.ToTensor()]),
download=True
)
elif dataset == 'lsun-bedroom':
res = 64
train_dir = "/home/ubuntu/code/sangyoon/forward-blur/datasets/bedroom_train/bedroom_train_lmdb/imgs"
dataset_train = dsets.LSUN(root='../forward-blur/datasets/bedroom_train',
classes=['bedroom_train'],
transform=transforms.Compose([transforms.Resize((res,res)), transforms.ToTensor()])
)
dataset_test = dsets.LSUN(root='../forward-blur/datasets/bedroom_train',
classes=['bedroom_train'],
transform=transforms.Compose([transforms.Resize((res,res)), transforms.ToTensor()])
)
elif dataset == 'lsun-church':
res = 64
train_dir = "/home/ubuntu/code/sangyoon/datasets/church_train/church_outdoor_train_lmdb/imgs"
dataset_train = dsets.LSUN(root='/home/ubuntu/code/sangyoon/datasets/church_train',
classes=['church_outdoor_train'],
transform=transforms.Compose([transforms.Resize((res,res)), transforms.ToTensor()])
)
dataset_test = dsets.LSUN(root='/home/ubuntu/code/sangyoon/datasets/church_test',
classes=['church_outdoor_val'],
transform=transforms.Compose([transforms.Resize((res,res)), transforms.ToTensor()])
)
fid_eval = FID(real_dir = train_dir, device = device)
resolution = dataset_train[0][0].shape[-1]
input_nc = dataset_train[0][0].shape[0]
ksize = resolution * 2 - 1
pad = 0
# define forward blur
kernel = gaussian_kernel_1d(ksize, sig)
blur = Deblurring(kernel, input_nc, resolution, device=device)
print("blur.U_small.shape:", blur.U_small.shape)
D_diag = blur.singulars()
fb = ForwardBlurIncreasing(N=N, beta_min=beta_min, beta_max=beta_max, sig=sig, sig_max = opt.sig_max, sig_min = opt.sig_min, D_diag=D_diag,
blur=blur, channel=input_nc, device=device, noise_schedule=opt.noise_schedule, resolution=resolution, pad=pad, f_type=opt.f_type)
dir = os.path.join('experiments', opt.name)
writer = SummaryWriter(dir)
# # Bedroom config - ADM
# model_channels = 256
# num_res_blocks = 2
# channel_mult = (1, 1, 2, 2, 4, 4)
# attention_resolutions = [8,16,32]
# out_channels = 6 # mean and varaince
# dropout = 0.1
# resblock_updown = True
# # bsize 256
# # iter -- cat (200K), horse (250K), bedroom (500K)
# lr = 1e-5 # Fine tuning
# model = UNetModel(image_size = resolution, in_channels = input_nc, model_channels = model_channels, out_channels = out_channels,
# num_res_blocks=num_res_blocks, channel_mult=channel_mult, attention_resolutions=attention_resolutions,
# dropout = dropout, resblock_updown = resblock_updown)
model = UNetModel(resolution, input_nc, 128, input_nc, blur = blur, dropout=opt.dropout, freq_feat = opt.freq_feat)
if not opt.ckpt == '' and os.path.exists(opt.ckpt):
model.load_state_dict(torch.load(opt.ckpt))
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
# dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
model = DataParallel(model)
model.to(device)
print("input_nc", input_nc, "resolution", resolution)
data_loader = torch.utils.data.DataLoader(dataset=dataset_train,
batch_size=bsize,
shuffle=True,
drop_last=True)
data_loader_test = torch.utils.data.DataLoader(dataset=dataset_test,
batch_size=bsize,
shuffle=False,
drop_last=True)
optimizer = optim.Adam(model.parameters(), lr=opt.lr)
if opt.use_ema:
optimizer = EMA(optimizer, ema_decay=opt.ema_decay)
# forward process visualization
sample = dataset_train[1][0].unsqueeze(0)
x_0 = sample[:4]
x_0 = x_0.to(device)
i = np.array([500] * x_0.shape[0])
i = torch.from_numpy(i).to(device)
fb.sanity(x_0, i)
sample_list = []
for i in range(0, N+1, N//10):
if i == 0:
sample_list.append(x_0[0])
continue
i = np.array([i] * x_0.shape[0])
i = torch.from_numpy(i).to(device)
x_i = fb.get_x_i(x_0, i)
sample_list.append(x_i[0])
print(f"x_{i.item()}.std() = {x_i.std()}")
print(f"x_{i.item()}.mean() = {x_i.mean()}")
grid_sample = torch.cat(sample_list, dim=2)
utils.tensor_imsave(grid_sample, "./" + dir, "forward_process.jpg")
with open(os.path.join(dir, "config.json"), "w") as json_file:
json.dump(vars(opt), json_file)
import time
meta_iter = 0
for step in range(opt.max_iter):
if not opt.inference:
elips = time.time()
try:
x_0, _ = train_iter.next()
except:
train_iter = iter(data_loader)
image, _ = train_iter.next()
"""
training
"""
assert x_0.shape[-1] == resolution, f"{x_0.shape}"
i = np.random.uniform(1 / N, 1, size = (x_0.shape[0])) * N
i = torch.from_numpy(i).to(device).type(torch.long)
x_0 = x_0.to(device)
x_i, eps = fb.get_x_i(x_0, i, return_eps = True)
if opt.loss_type == "sm_simple":
loss = fb.get_loss_i_simple(model, x_0, x_i, i)
elif opt.loss_type == "eps_simple":
loss = fb.get_loss_i_eps_simple(model, x_i, i, eps)
elif opt.loss_type == "sm_exact":
loss = fb.get_loss_i_exact(model, x_0, x_i, i)
elif opt.loss_type == "std_matching":
loss = fb.get_loss_i_std_matching(model, x_i, i, eps)
writer.add_scalar('loss_train', loss, step)
loss.backward()
optimizer.step()
optimizer.zero_grad()
if step % 100 == 0:
print(step, loss)
# print(f"time: {time.time() - elips}")
# Calcuate FID
if step > 240001:
fid_iter = opt.fid_iter
else:
fid_iter = 240000
if (step % fid_iter == 0 and step > 0):
id = 0
if not os.path.exists(os.path.join("./",dir, f"{step}")):
os.mkdir(os.path.join("./",dir, f"{step}"))
with torch.no_grad():
if opt.use_ema:
optimizer.swap_parameters_with_ema(store_params_in_ema=True)
model = model.eval()
for _ in range(opt.fid_num_samples // opt.fid_bsize):
i = np.array([opt.N - 1] * opt.fid_bsize)
i = torch.from_numpy(i).to(device)
pred = fb.get_x_N([opt.fid_bsize, input_nc, resolution, resolution], i)
for i in reversed(range(1, opt.N)):
i = np.array([i] * opt.fid_bsize)
i = torch.from_numpy(i).to(device)
if opt.loss_type == "sm_simple":
s = model(pred, i)
elif opt.loss_type == "eps_simple":
eps = model(pred, i)
s = fb.get_score_from_eps(eps, i)
elif opt.loss_type == "sm_exact":
s = model(pred, i)
elif opt.loss_type == "std_matching":
std = model(pred, i)
s = fb.get_score_from_std(std, i)
else:
raise NotImplementedError
s = fb.U_I_minus_B_Ut(s, i)
rms = lambda x: torch.sqrt(torch.mean(x ** 2))
# print(f"rms(s) * fb._beta_i(i) = {rms(s) * fb._beta_i(i)[0]}")
hf = pred - fb.W(pred, i)
# Anderson theorem
pred1 = pred + hf # unsharpening mask filtering
pred2 = pred1 + s # # denoising
if i[0] > 2:
pred = pred2 + fb.U_I_minus_B_sqrt_Ut(torch.randn_like(pred), i) # inject noise
else:
pred = pred2
# print(f"i = {i[0]}, rmse = {torch.sqrt(torch.mean(pred**2))}, mean = {torch.mean(pred)} std = {torch.std(pred)}" )
for sample in pred:
save_image(sample, os.path.join(dir, f"{step}", f"{id:05d}.png"))
id += 1
if opt.use_ema:
optimizer.swap_parameters_with_ema(store_params_in_ema=True)
model = model.train()
fid = fid_eval(os.path.join(dir, f"{step}"))
writer.add_scalar('fid', fid, step)
print(f"step {step}, fid = {fid}")
if (step % opt.eval_iter == 0 and step > 0) or opt.inference:
"""
sampling (eval)
"""
cnt = 0
loss = 0
with torch.no_grad():
if opt.use_ema:
optimizer.swap_parameters_with_ema(store_params_in_ema=True)
model = model.eval()
if opt.ode:
raise NotImplementedError
def to_flattened_numpy(x):
"""Flatten a torch tensor `x` and convert it to numpy."""
return x.detach().cpu().numpy().reshape((-1,))
def from_flattened_numpy(x, shape):
"""Form a torch tensor with the given `shape` from a flattened numpy array `x`."""
return torch.from_numpy(x.reshape(shape))
def ode_func(i, y):
i = int(i*N)
print(f"i = {i}")
y = from_flattened_numpy(y, [bsize, input_nc, resolution, resolution]).to(device).type(torch.float32)
i = np.array([N - 1] * bsize)
i = torch.from_numpy(i).to(device)
if opt.loss_type == "sm_simple":
s = model(y, i)
elif opt.loss_type == "eps_simple":
eps = model(y, i)
s = fb.get_score_from_eps(eps, i)
elif opt.loss_type == "sm_exact":
s = model(y, i)
elif opt.loss_type == "std_matching":
std = model(y, i)
s = fb.get_score_from_std(std, i)
else:
raise NotImplementedError
s = fb.U_I_minus_B_Ut(s, i)
hf = y - fb.W(y, i)
dt = - 1.0 / N
drift = (s/2 + hf) / dt
drift = to_flattened_numpy(drift)
return drift
x_N = fb.get_x_N([bsize, input_nc, resolution, resolution], N)
solution = solve_ivp(ode_func, (1, 1e-3), to_flattened_numpy(x_N),
rtol=1e-3, atol=1e-3, method="RK45")
nfe = solution.nfev
solution = torch.tensor(solution.y[:, -1])
gitextract_8zm764cm/ ├── EMA.py ├── LICENSE ├── README.md ├── blur_diffusion.py ├── eval_x0hat.py ├── eval_x0hat.sh ├── fid.py ├── guided_diffusion/ │ ├── __init__.py │ ├── dist_util.py │ ├── fp16_util.py │ ├── gaussian_diffusion.py │ ├── image_datasets.py │ ├── logger.py │ ├── losses.py │ ├── nn.py │ ├── resample.py │ ├── respace.py │ ├── script_util.py │ ├── train_util.py │ └── unet.py ├── main.py ├── pytorch_fid/ │ ├── __init__.py │ ├── __main__.py │ ├── fid_score.py │ └── inception.py ├── train.sh └── utils.py
SYMBOL INDEX (364 symbols across 19 files)
FILE: EMA.py
class EMA (line 17) | class EMA(Optimizer):
method __init__ (line 18) | def __init__(self, opt, ema_decay):
method step (line 25) | def step(self, *args, **kwargs):
method load_state_dict (line 64) | def load_state_dict(self, state_dict):
method swap_parameters_with_ema (line 71) | def swap_parameters_with_ema(self, store_params_in_ema):
FILE: blur_diffusion.py
function gaussian_kernel_1d (line 10) | def gaussian_kernel_1d(kernel_size, sigma):
function betas_for_alpha_bar (line 22) | def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.9...
class ExpSchedule (line 34) | class ExpSchedule:
method __init__ (line 35) | def __init__(self, N, offset = 1e-4):
method get_betas (line 45) | def get_betas(self):
class ForwardBlurIncreasing (line 51) | class ForwardBlurIncreasing:
method __init__ (line 52) | def __init__(self, N, beta_min, beta_max, sig, sig_min, sig_max, D_dia...
method _beta_i (line 184) | def _beta_i(self, i):
method _alpha_i (line 187) | def _alpha_i(self, i):
method _alphas_bar_i (line 190) | def _alphas_bar_i(self, i):
method _Bhat_i (line 193) | def _Bhat_i(self, i):
method _Bhat_sq_i (line 197) | def _Bhat_sq_i(self, i):
method get_mean (line 204) | def get_mean(self, x0, i):
method get_std (line 210) | def get_std(self,i, noise):
method get_var (line 216) | def get_var(self, x0, i, noise=None):
method W (line 223) | def W(self, x, i):
method W_inv (line 233) | def W_inv(self, x, i):
method U_I_minus_B_Ut (line 241) | def U_I_minus_B_Ut(self, x, i): # U(I-B)UT
method U_I_minus_B_sqrt_Ut (line 250) | def U_I_minus_B_sqrt_Ut(self, x, i): # U(sqrt(I-B))UT
method get_x_i (line 260) | def get_x_i(self, x0, i, return_eps = False):
method get_x_N (line 284) | def get_x_N(self, x0_shape, N):
method get_x0_from_eps (line 308) | def get_x0_from_eps(self, xi, eps, i):
method get_score_gt (line 323) | def get_score_gt(self, xi, x0, i):
method get_score_gt2 (line 341) | def get_score_gt2(self, xi, x0, i):
method sanity (line 360) | def sanity(self, x_0, i):
method get_score_from_eps (line 402) | def get_score_from_eps(self, eps, i):
method get_score_from_std (line 409) | def get_score_from_std(self, std, i):
method get_loss_i_exact (line 415) | def get_loss_i_exact(self, model, x0, xi, i):
method get_loss_i_simple (line 423) | def get_loss_i_simple(self, model, x0, xi, i):
method get_loss_i_eps_simple (line 444) | def get_loss_i_eps_simple(self, model, x_i, i, eps):
method get_loss_i_std_matching (line 448) | def get_loss_i_std_matching(self, model, x_i, i, eps):
class H_functions (line 458) | class H_functions:
method V (line 466) | def V(self, vec):
method Vt (line 472) | def Vt(self, vec):
method U (line 478) | def U(self, vec):
method Ut (line 484) | def Ut(self, vec):
method singulars (line 490) | def singulars(self):
method add_zeros (line 496) | def add_zeros(self, vec):
method H (line 502) | def H(self, vec):
method Ht (line 510) | def Ht(self, vec):
method H_pinv (line 518) | def H_pinv(self, vec):
class Deblurring (line 528) | class Deblurring(H_functions):
method mat_by_img (line 529) | def mat_by_img(self, M, v):
method img_by_mat (line 533) | def img_by_mat(self, v, M):
method __init__ (line 537) | def __init__(self, kernel, channels, img_dim, device):
method V (line 564) | def V(self, vec):
method Vt (line 574) | def Vt(self, vec):
method U (line 582) | def U(self, vec):
method Ut (line 592) | def Ut(self, vec):
method conv1d_col_matmul (line 600) | def conv1d_col_matmul(self, x):
method conv1d_row_matmul (line 603) | def conv1d_row_matmul(self, x):
method conv2d_sep_matmul (line 606) | def conv2d_sep_matmul(self, x):
method singulars (line 609) | def singulars(self):
method update_singulars (line 615) | def update_singulars(self, new_singulars):
method add_zeros (line 618) | def add_zeros(self, vec):
FILE: fid.py
class FID (line 4) | class FID(object):
method __init__ (line 5) | def __init__(self, real_dir, device, bsize = 128):
method __call__ (line 20) | def __call__(self, fake_dir):
FILE: guided_diffusion/dist_util.py
function setup_dist (line 21) | def setup_dist():
function dev (line 51) | def dev():
function load_state_dict (line 60) | def load_state_dict(path, **kwargs):
function sync_params (line 83) | def sync_params(params):
function _find_free_port (line 92) | def _find_free_port():
FILE: guided_diffusion/fp16_util.py
function convert_module_to_f16 (line 15) | def convert_module_to_f16(l):
function convert_module_to_f32 (line 25) | def convert_module_to_f32(l):
function make_master_params (line 35) | def make_master_params(param_groups_and_shapes):
function model_grads_to_master_grads (line 52) | def model_grads_to_master_grads(param_groups_and_shapes, master_params):
function master_params_to_model_params (line 65) | def master_params_to_model_params(param_groups_and_shapes, master_params):
function unflatten_master_params (line 78) | def unflatten_master_params(param_group, master_param):
function get_param_groups_and_shapes (line 82) | def get_param_groups_and_shapes(named_model_params):
function master_params_to_state_dict (line 95) | def master_params_to_state_dict(
function state_dict_to_master_params (line 116) | def state_dict_to_master_params(model, state_dict, use_fp16):
function zero_master_grads (line 128) | def zero_master_grads(master_params):
function zero_grad (line 133) | def zero_grad(model_params):
function param_grad_or_zeros (line 141) | def param_grad_or_zeros(param):
class MixedPrecisionTrainer (line 148) | class MixedPrecisionTrainer:
method __init__ (line 149) | def __init__(
method zero_grad (line 173) | def zero_grad(self):
method backward (line 176) | def backward(self, loss: th.Tensor):
method optimize (line 183) | def optimize(self, opt: th.optim.Optimizer):
method _optimize_fp16 (line 189) | def _optimize_fp16(self, opt: th.optim.Optimizer):
method _optimize_normal (line 209) | def _optimize_normal(self, opt: th.optim.Optimizer):
method _compute_norms (line 216) | def _compute_norms(self, grad_scale=1.0):
method master_params_to_state_dict (line 226) | def master_params_to_state_dict(self, master_params):
method state_dict_to_master_params (line 231) | def state_dict_to_master_params(self, state_dict):
function check_overflow (line 235) | def check_overflow(value):
FILE: guided_diffusion/gaussian_diffusion.py
function get_named_beta_schedule (line 18) | def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
function betas_for_alpha_bar (line 45) | def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.9...
class ModelMeanType (line 65) | class ModelMeanType(enum.Enum):
class ModelVarType (line 75) | class ModelVarType(enum.Enum):
class LossType (line 89) | class LossType(enum.Enum):
method is_vb (line 97) | def is_vb(self):
class GaussianDiffusion (line 101) | class GaussianDiffusion:
method __init__ (line 118) | def __init__(
method q_mean_variance (line 171) | def q_mean_variance(self, x_start, t):
method q_sample (line 188) | def q_sample(self, x_start, t, noise=None):
method q_sample_blur (line 208) | def q_sample_blur(self, x_start, t, noise=None):
method q_posterior_mean_variance (line 226) | def q_posterior_mean_variance(self, x_start, x_t, t):
method p_mean_variance (line 250) | def p_mean_variance(
method _predict_xstart_from_eps (line 346) | def _predict_xstart_from_eps(self, x_t, t, eps):
method _predict_xstart_from_xprev (line 353) | def _predict_xstart_from_xprev(self, x_t, t, xprev):
method _predict_eps_from_xstart (line 363) | def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
method _scale_timesteps (line 369) | def _scale_timesteps(self, t):
method condition_mean (line 374) | def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
method condition_score (line 389) | def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
method p_sample (line 412) | def p_sample(
method p_sample_blur (line 458) | def p_sample_blur(
method p_sample_loop (line 506) | def p_sample_loop(
method p_sample_loop_progressive (line 557) | def p_sample_loop_progressive(
method ddim_sample (line 610) | def ddim_sample(
method ddim_reverse_sample (line 660) | def ddim_reverse_sample(
method ddim_sample_loop (line 698) | def ddim_sample_loop(
method ddim_sample_loop_progressive (line 732) | def ddim_sample_loop_progressive(
method _vb_terms_bpd (line 782) | def _vb_terms_bpd(
method training_losses (line 817) | def training_losses(self, model, x_start, t, model_kwargs=None, noise=...
method _prior_bpd (line 897) | def _prior_bpd(self, x_start):
method calc_bpd_loop (line 915) | def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwar...
function _extract_into_tensor (line 998) | def _extract_into_tensor(arr, timesteps, broadcast_shape):
function tensor2pil (line 1013) | def tensor2pil(t):
function denorm (line 1016) | def denorm(t):
function tensor_imsave (line 1019) | def tensor_imsave(t, path, fname, denormalization=True, prt=False):
function pil_loader (line 1028) | def pil_loader(path):
function check_folder (line 1033) | def check_folder(log_dir):
FILE: guided_diffusion/image_datasets.py
function load_data (line 11) | def load_data(
function _list_image_files_recursively (line 70) | def _list_image_files_recursively(data_dir):
class ImageDataset (line 82) | class ImageDataset(Dataset):
method __init__ (line 83) | def __init__(
method __len__ (line 100) | def __len__(self):
method __getitem__ (line 103) | def __getitem__(self, idx):
function center_crop_arr (line 126) | def center_crop_arr(pil_image, image_size):
function random_crop_arr (line 146) | def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_f...
FILE: guided_diffusion/logger.py
class KVWriter (line 26) | class KVWriter(object):
method writekvs (line 27) | def writekvs(self, kvs):
class SeqWriter (line 31) | class SeqWriter(object):
method writeseq (line 32) | def writeseq(self, seq):
class HumanOutputFormat (line 36) | class HumanOutputFormat(KVWriter, SeqWriter):
method __init__ (line 37) | def __init__(self, filename_or_file):
method writekvs (line 48) | def writekvs(self, kvs):
method _truncate (line 80) | def _truncate(self, s):
method writeseq (line 84) | def writeseq(self, seq):
method close (line 93) | def close(self):
class JSONOutputFormat (line 98) | class JSONOutputFormat(KVWriter):
method __init__ (line 99) | def __init__(self, filename):
method writekvs (line 102) | def writekvs(self, kvs):
method close (line 109) | def close(self):
class CSVOutputFormat (line 113) | class CSVOutputFormat(KVWriter):
method __init__ (line 114) | def __init__(self, filename):
method writekvs (line 119) | def writekvs(self, kvs):
method close (line 146) | def close(self):
class TensorBoardOutputFormat (line 150) | class TensorBoardOutputFormat(KVWriter):
method __init__ (line 155) | def __init__(self, dir):
method writekvs (line 171) | def writekvs(self, kvs):
method close (line 185) | def close(self):
function make_output_format (line 191) | def make_output_format(format, ev_dir, log_suffix=""):
function logkv (line 212) | def logkv(key, val):
function logkv_mean (line 221) | def logkv_mean(key, val):
function logkvs (line 228) | def logkvs(d):
function dumpkvs (line 236) | def dumpkvs():
function getkvs (line 243) | def getkvs():
function log (line 247) | def log(*args, level=INFO):
function debug (line 254) | def debug(*args):
function info (line 258) | def info(*args):
function warn (line 262) | def warn(*args):
function error (line 266) | def error(*args):
function set_level (line 270) | def set_level(level):
function set_comm (line 277) | def set_comm(comm):
function get_dir (line 281) | def get_dir():
function profile_kv (line 294) | def profile_kv(scopename):
function profile (line 303) | def profile(n):
function get_current (line 325) | def get_current():
class Logger (line 332) | class Logger(object):
method __init__ (line 337) | def __init__(self, dir, output_formats, comm=None):
method logkv (line 347) | def logkv(self, key, val):
method logkv_mean (line 350) | def logkv_mean(self, key, val):
method dumpkvs (line 355) | def dumpkvs(self):
method log (line 376) | def log(self, *args, level=INFO):
method set_level (line 382) | def set_level(self, level):
method set_comm (line 385) | def set_comm(self, comm):
method get_dir (line 388) | def get_dir(self):
method close (line 391) | def close(self):
method _do_log (line 397) | def _do_log(self, args):
function get_rank_without_mpi_import (line 403) | def get_rank_without_mpi_import():
function mpi_weighted_mean (line 412) | def mpi_weighted_mean(comm, local_name2valcount):
function configure (line 442) | def configure(dir=None, format_strs=None, comm=None, log_suffix=""):
function _configure_default_logger (line 475) | def _configure_default_logger():
function reset (line 480) | def reset():
function scoped_configure (line 488) | def scoped_configure(dir=None, format_strs=None, comm=None):
FILE: guided_diffusion/losses.py
function normal_kl (line 12) | def normal_kl(mean1, logvar1, mean2, logvar2):
function approx_standard_normal_cdf (line 42) | def approx_standard_normal_cdf(x):
function discretized_gaussian_log_likelihood (line 50) | def discretized_gaussian_log_likelihood(x, *, means, log_scales):
FILE: guided_diffusion/nn.py
class SiLU (line 12) | class SiLU(nn.Module):
method forward (line 13) | def forward(self, x):
class GroupNorm32 (line 17) | class GroupNorm32(nn.GroupNorm):
method forward (line 18) | def forward(self, x):
function conv_nd (line 22) | def conv_nd(dims, *args, **kwargs):
function linear (line 35) | def linear(*args, **kwargs):
function avg_pool_nd (line 42) | def avg_pool_nd(dims, *args, **kwargs):
function update_ema (line 55) | def update_ema(target_params, source_params, rate=0.99):
function zero_module (line 68) | def zero_module(module):
function scale_module (line 77) | def scale_module(module, scale):
function mean_flat (line 86) | def mean_flat(tensor):
function normalization (line 93) | def normalization(channels):
function timestep_embedding (line 103) | def timestep_embedding(timesteps, dim, max_period=10000):
function checkpoint (line 124) | def checkpoint(func, inputs, params, flag):
class CheckpointFunction (line 142) | class CheckpointFunction(th.autograd.Function):
method forward (line 144) | def forward(ctx, run_function, length, *args):
method backward (line 153) | def backward(ctx, *output_grads):
FILE: guided_diffusion/resample.py
function create_named_schedule_sampler (line 8) | def create_named_schedule_sampler(name, diffusion):
class ScheduleSampler (line 23) | class ScheduleSampler(ABC):
method weights (line 35) | def weights(self):
method sample (line 42) | def sample(self, batch_size, device):
class UniformSampler (line 64) | class UniformSampler(ScheduleSampler):
method __init__ (line 65) | def __init__(self, diffusion):
method weights (line 69) | def weights(self):
class LossAwareSampler (line 73) | class LossAwareSampler(ScheduleSampler):
method update_with_local_losses (line 74) | def update_with_local_losses(self, local_ts, local_losses):
method update_with_all_losses (line 110) | def update_with_all_losses(self, ts, losses):
class LossSecondMomentResampler (line 127) | class LossSecondMomentResampler(LossAwareSampler):
method __init__ (line 128) | def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001):
method weights (line 138) | def weights(self):
method update_with_all_losses (line 147) | def update_with_all_losses(self, ts, losses):
method _warmed_up (line 157) | def _warmed_up(self):
FILE: guided_diffusion/respace.py
function space_timesteps (line 7) | def space_timesteps(num_timesteps, section_counts):
class SpacedDiffusion (line 63) | class SpacedDiffusion(GaussianDiffusion):
method __init__ (line 72) | def __init__(self, use_timesteps, **kwargs):
method p_mean_variance (line 88) | def p_mean_variance(
method training_losses (line 93) | def training_losses(
method condition_mean (line 98) | def condition_mean(self, cond_fn, *args, **kwargs):
method condition_score (line 101) | def condition_score(self, cond_fn, *args, **kwargs):
method _wrap_model (line 104) | def _wrap_model(self, model):
method _scale_timesteps (line 111) | def _scale_timesteps(self, t):
class _WrappedModel (line 116) | class _WrappedModel:
method __init__ (line 117) | def __init__(self, model, timestep_map, rescale_timesteps, original_nu...
method __call__ (line 123) | def __call__(self, x, ts, **kwargs):
FILE: guided_diffusion/script_util.py
function diffusion_defaults (line 11) | def diffusion_defaults():
function classifier_defaults (line 27) | def classifier_defaults():
function model_and_diffusion_defaults (line 43) | def model_and_diffusion_defaults():
function classifier_and_diffusion_defaults (line 68) | def classifier_and_diffusion_defaults():
function create_model_and_diffusion (line 74) | def create_model_and_diffusion(
function create_model (line 130) | def create_model(
function create_classifier_and_diffusion (line 187) | def create_classifier_and_diffusion(
function create_classifier (line 228) | def create_classifier(
function sr_model_and_diffusion_defaults (line 269) | def sr_model_and_diffusion_defaults():
function sr_create_model_and_diffusion (line 280) | def sr_create_model_and_diffusion(
function sr_create_model (line 334) | def sr_create_model(
function create_gaussian_diffusion (line 386) | def create_gaussian_diffusion(
function add_dict_to_argparser (line 427) | def add_dict_to_argparser(parser, default_dict):
function args_to_dict (line 437) | def args_to_dict(args, keys):
function str2bool (line 441) | def str2bool(v):
FILE: guided_diffusion/train_util.py
class TrainLoop (line 22) | class TrainLoop:
method __init__ (line 23) | def __init__(
method _load_and_sync_parameters (line 110) | def _load_and_sync_parameters(self):
method _load_ema_parameters (line 125) | def _load_ema_parameters(self, rate):
method _load_optimizer_state (line 141) | def _load_optimizer_state(self):
method run_loop (line 153) | def run_loop(self):
method run_step (line 172) | def run_step(self, batch, cond):
method forward_backward (line 180) | def forward_backward(self, batch, cond):
method _update_ema (line 217) | def _update_ema(self):
method _anneal_lr (line 221) | def _anneal_lr(self):
method log_step (line 229) | def log_step(self):
method save (line 233) | def save(self):
function parse_resume_step_from_filename (line 259) | def parse_resume_step_from_filename(filename):
function get_blob_logdir (line 274) | def get_blob_logdir():
function find_resume_checkpoint (line 280) | def find_resume_checkpoint():
function find_ema_checkpoint (line 286) | def find_ema_checkpoint(main_checkpoint, step, rate):
function log_loss_dict (line 296) | def log_loss_dict(diffusion, ts, losses):
FILE: guided_diffusion/unet.py
class AttentionPool2d (line 22) | class AttentionPool2d(nn.Module):
method __init__ (line 27) | def __init__(
method forward (line 43) | def forward(self, x):
class TimestepBlock (line 54) | class TimestepBlock(nn.Module):
method forward (line 60) | def forward(self, x, emb):
class TimestepEmbedSequential (line 66) | class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
method forward (line 72) | def forward(self, x, emb):
class Upsample (line 81) | class Upsample(nn.Module):
method __init__ (line 91) | def __init__(self, channels, use_conv, dims=2, out_channels=None):
method forward (line 100) | def forward(self, x):
class Downsample (line 113) | class Downsample(nn.Module):
method __init__ (line 123) | def __init__(self, channels, use_conv, dims=2, out_channels=None):
method forward (line 138) | def forward(self, x):
class ResBlock (line 143) | class ResBlock(TimestepBlock):
method __init__ (line 160) | def __init__(
method forward (line 224) | def forward(self, x, emb):
method _forward (line 236) | def _forward(self, x, emb):
class AttentionBlock (line 259) | class AttentionBlock(nn.Module):
method __init__ (line 267) | def __init__(
method forward (line 296) | def forward(self, x):
method _forward (line 299) | def _forward(self, x):
function count_flops_attn (line 308) | def count_flops_attn(model, _x, y):
class QKVAttentionLegacy (line 328) | class QKVAttentionLegacy(nn.Module):
method __init__ (line 333) | def __init__(self, n_heads):
method forward (line 337) | def forward(self, qkv):
method count_flops (line 357) | def count_flops(model, _x, y):
class QKVAttention (line 361) | class QKVAttention(nn.Module):
method __init__ (line 366) | def __init__(self, n_heads):
method forward (line 370) | def forward(self, qkv):
method count_flops (line 392) | def count_flops(model, _x, y):
class UNetModel (line 396) | class UNetModel(nn.Module):
method __init__ (line 427) | def __init__(
method convert_to_fp16 (line 631) | def convert_to_fp16(self):
method convert_to_fp32 (line 639) | def convert_to_fp32(self):
method forward (line 647) | def forward(self, x, timesteps, y=None):
class SuperResModel (line 684) | class SuperResModel(UNetModel):
method __init__ (line 691) | def __init__(self, image_size, in_channels, *args, **kwargs):
method forward (line 694) | def forward(self, x, timesteps, low_res=None, **kwargs):
class EncoderUNetModel (line 701) | class EncoderUNetModel(nn.Module):
method __init__ (line 708) | def __init__(
method convert_to_fp16 (line 875) | def convert_to_fp16(self):
method convert_to_fp32 (line 882) | def convert_to_fp32(self):
method forward (line 889) | def forward(self, x, timesteps):
FILE: main.py
function to_flattened_numpy (line 306) | def to_flattened_numpy(x):
function from_flattened_numpy (line 309) | def from_flattened_numpy(x, shape):
function ode_func (line 312) | def ode_func(i, y):
FILE: pytorch_fid/fid_score.py
function tqdm (line 49) | def tqdm(x):
class ImagePathDataset (line 77) | class ImagePathDataset(torch.utils.data.Dataset):
method __init__ (line 78) | def __init__(self, files, transforms=None):
method __len__ (line 82) | def __len__(self):
method __getitem__ (line 85) | def __getitem__(self, i):
function get_activations (line 93) | def get_activations(files, model, batch_size=50, dims=2048, device='cpu',
function calculate_frechet_distance (line 152) | def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
function calculate_activation_statistics (line 209) | def calculate_activation_statistics(files, model, batch_size=50, dims=2048,
function compute_statistics_of_path (line 234) | def compute_statistics_of_path(path, model, batch_size, dims, device,
function calculate_fid_given_paths (line 251) | def calculate_fid_given_paths(paths, batch_size, device, dims, num_worke...
function main (line 270) | def main():
FILE: pytorch_fid/inception.py
class InceptionV3 (line 16) | class InceptionV3(nn.Module):
method __init__ (line 31) | def __init__(self,
method forward (line 129) | def forward(self, inp):
function _inception_v3 (line 166) | def _inception_v3(*args, **kwargs):
function fid_inception_v3 (line 184) | def fid_inception_v3():
class FIDInceptionA (line 211) | class FIDInceptionA(torchvision.models.inception.InceptionA):
method __init__ (line 213) | def __init__(self, in_channels, pool_features):
method forward (line 216) | def forward(self, x):
class FIDInceptionC (line 236) | class FIDInceptionC(torchvision.models.inception.InceptionC):
method __init__ (line 238) | def __init__(self, in_channels, channels_7x7):
method forward (line 241) | def forward(self, x):
class FIDInceptionE_1 (line 264) | class FIDInceptionE_1(torchvision.models.inception.InceptionE):
method __init__ (line 266) | def __init__(self, in_channels):
method forward (line 269) | def forward(self, x):
class FIDInceptionE_2 (line 297) | class FIDInceptionE_2(torchvision.models.inception.InceptionE):
method __init__ (line 299) | def __init__(self, in_channels):
method forward (line 302) | def forward(self, x):
FILE: utils.py
function tensor2pil (line 9) | def tensor2pil(t):
function denorm (line 12) | def denorm(t):
function tensor_imsave (line 14) | def tensor_imsave(t, path, fname, denormalization=False, prt=False):
function pil_loader (line 23) | def pil_loader(path):
function check_folder (line 28) | def check_folder(log_dir):
function freeze_model (line 32) | def freeze_model(model):
class TVLoss (line 35) | class TVLoss(nn.Module):
method __init__ (line 36) | def __init__(self,TVLoss_weight=1):
method forward (line 40) | def forward(self,x):
method _tensor_size (line 50) | def _tensor_size(self,t):
function pil_batch_loader (line 53) | def pil_batch_loader(root_path, bsize):
class PSNR (line 63) | class PSNR(object):
method __init__ (line 64) | def __init__(self, gpu, val_max=1, val_min=0, ycbcr=True):
method __call__ (line 71) | def __call__(self,x,y):
function rgb_to_ycbcr (line 92) | def rgb_to_ycbcr(image):
function clear (line 118) | def clear(x):
function normalize_np (line 121) | def normalize_np(x):
function pad (line 126) | def pad(x, pad):
function crop (line 129) | def crop(x, pad):
Condensed preview — 27 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (239K chars).
[
{
"path": "EMA.py",
"chars": 3321,
"preview": "# ---------------------------------------------------------------\n# Copyright (c) 2022, NVIDIA CORPORATION. All rights r"
},
{
"path": "LICENSE",
"chars": 1068,
"preview": "MIT License\n\nCopyright (c) 2022 Sangyun Lee\n\nPermission is hereby granted, free of charge, to any person obtaining a cop"
},
{
"path": "README.md",
"chars": 1875,
"preview": "# blur-diffusion\nThis is the codebase for our paper Progressive Deblurring of Diffusion Models for Coarse-to-Fine Image "
},
{
"path": "blur_diffusion.py",
"chars": 24514,
"preview": "import torch\nimport torch.nn.functional as F\nimport time\nimport utils\nimport os\nimport numpy as np\n\nfrom utils import cl"
},
{
"path": "eval_x0hat.py",
"chars": 5155,
"preview": "import torch\nfrom PIL import Image\nfrom collections import defaultdict\nimport torch.nn.functional as TF\nimport torchvisi"
},
{
"path": "eval_x0hat.sh",
"chars": 379,
"preview": "#!/bin/bash\nexport OMP_NUM_THREADS=1\nexport CUDA_VISIBLE_DEVICES=1\n\npython eval_x0hat.py \\\n--name test \\\n--noise_schedul"
},
{
"path": "fid.py",
"chars": 1207,
"preview": "from pytorch_fid.fid_score import compute_statistics_of_path, calculate_frechet_distance\nfrom pytorch_fid.inception impo"
},
{
"path": "guided_diffusion/__init__.py",
"chars": 74,
"preview": "\"\"\"\nCodebase for \"Improved Denoising Diffusion Probabilistic Models\".\n\"\"\"\n"
},
{
"path": "guided_diffusion/dist_util.py",
"chars": 2751,
"preview": "\"\"\"\nHelpers for distributed training.\n\"\"\"\n\nimport io\nimport os\nimport socket\n\nimport blobfile as bf\nfrom mpi4py import M"
},
{
"path": "guided_diffusion/fp16_util.py",
"chars": 7941,
"preview": "\"\"\"\nHelpers to train with 16-bit precision.\n\"\"\"\n\nimport numpy as np\nimport torch as th\nimport torch.nn as nn\nfrom torch."
},
{
"path": "guided_diffusion/gaussian_diffusion.py",
"chars": 39274,
"preview": "\"\"\"\nThis code started out as a PyTorch port of Ho et al's diffusion models:\nhttps://github.com/hojonathanho/diffusion/bl"
},
{
"path": "guided_diffusion/image_datasets.py",
"chars": 5930,
"preview": "import math\nimport random\n\nfrom PIL import Image\nimport blobfile as bf\nfrom mpi4py import MPI\nimport numpy as np\nfrom to"
},
{
"path": "guided_diffusion/logger.py",
"chars": 14030,
"preview": "\"\"\"\nLogger copied from OpenAI baselines to avoid extra RL-based dependencies:\nhttps://github.com/openai/baselines/blob/e"
},
{
"path": "guided_diffusion/losses.py",
"chars": 2534,
"preview": "\"\"\"\nHelpers for various likelihood-based losses. These are ported from the original\nHo et al. diffusion models codebase:"
},
{
"path": "guided_diffusion/nn.py",
"chars": 5020,
"preview": "\"\"\"\nVarious utilities for neural networks.\n\"\"\"\n\nimport math\n\nimport torch as th\nimport torch.nn as nn\n\n\n# PyTorch 1.7 ha"
},
{
"path": "guided_diffusion/resample.py",
"chars": 5911,
"preview": "from abc import ABC, abstractmethod\n\nimport numpy as np\nimport torch as th\nimport torch.distributed as dist\n\n\ndef create"
},
{
"path": "guided_diffusion/respace.py",
"chars": 5193,
"preview": "import numpy as np\nimport torch as th\n\nfrom .gaussian_diffusion import GaussianDiffusion\n\n\ndef space_timesteps(num_times"
},
{
"path": "guided_diffusion/script_util.py",
"chars": 12349,
"preview": "import argparse\nimport inspect\n\nfrom . import gaussian_diffusion as gd\nfrom .respace import SpacedDiffusion, space_times"
},
{
"path": "guided_diffusion/train_util.py",
"chars": 10722,
"preview": "import copy\nimport functools\nimport os\n\nimport blobfile as bf\nimport torch as th\nimport torch.distributed as dist\nfrom t"
},
{
"path": "guided_diffusion/unet.py",
"chars": 32090,
"preview": "from abc import abstractmethod\n\nimport math\n\nimport numpy as np\nimport torch as th\nimport torch.nn as nn\nimport torch.nn"
},
{
"path": "main.py",
"chars": 20342,
"preview": "import torch\nfrom PIL import Image\nfrom collections import defaultdict\nimport torch.nn.functional as TF\nimport torchvisi"
},
{
"path": "pytorch_fid/__init__.py",
"chars": 111,
"preview": "\"\"\"\nhttps://github.com/mseitzer/pytorch-fid/blob/master/src/pytorch_fid/fid_score.py\n\"\"\"\n__version__ = '0.2.1'\n"
},
{
"path": "pytorch_fid/__main__.py",
"chars": 59,
"preview": "import pytorch_fid.fid_score\n\npytorch_fid.fid_score.main()\n"
},
{
"path": "pytorch_fid/fid_score.py",
"chars": 10837,
"preview": "\"\"\"Calculates the Frechet Inception Distance (FID) to evalulate GANs\n\nThe FID metric calculates the distance between two"
},
{
"path": "pytorch_fid/inception.py",
"chars": 12193,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\n\ntry:\n from torchvision.models."
},
{
"path": "train.sh",
"chars": 312,
"preview": "export OMP_NUM_THREADS=1\nexport CUDA_VISIBLE_DEVICES=0\n\npython main.py \\\n--name test \\\n--dataset lsun-bedroom \\\n--use_em"
},
{
"path": "utils.py",
"chars": 4319,
"preview": "from torchvision import transforms\nfrom PIL import Image\nimport os\nfrom torch import nn\nimport torch\nimport glob\nimport "
}
]
About this extraction
This page contains the full source code of the sangyun884/blur-diffusion GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 27 files (224.1 KB), approximately 55.4k tokens, and a symbol index with 364 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.