|
|
|
@@ -19,8 +19,6 @@ from tqdm import tqdm
|
|
|
|
|
from einops import rearrange, reduce
|
|
|
|
|
from einops.layers.torch import Rearrange
|
|
|
|
|
|
|
|
|
|
from ema_pytorch import EMA
|
|
|
|
|
|
|
|
|
|
# helpers functions
|
|
|
|
|
|
|
|
|
|
def exists(x):
|
|
|
|
@@ -52,6 +50,21 @@ def unnormalize_to_zero_to_one(t):
|
|
|
|
|
|
|
|
|
|
# small helper modules
|
|
|
|
|
|
|
|
|
|
class EMA():
|
|
|
|
|
def __init__(self, beta):
|
|
|
|
|
super().__init__()
|
|
|
|
|
self.beta = beta
|
|
|
|
|
|
|
|
|
|
def update_model_average(self, ma_model, current_model):
|
|
|
|
|
for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):
|
|
|
|
|
old_weight, up_weight = ma_params.data, current_params.data
|
|
|
|
|
ma_params.data = self.update_average(old_weight, up_weight)
|
|
|
|
|
|
|
|
|
|
def update_average(self, old, new):
|
|
|
|
|
if old is None:
|
|
|
|
|
return new
|
|
|
|
|
return old * self.beta + (1 - self.beta) * new
|
|
|
|
|
|
|
|
|
|
class Residual(nn.Module):
|
|
|
|
|
def __init__(self, fn):
|
|
|
|
|
super().__init__()
|
|
|
|
@@ -60,6 +73,20 @@ class Residual(nn.Module):
|
|
|
|
|
def forward(self, x, *args, **kwargs):
|
|
|
|
|
return self.fn(x, *args, **kwargs) + x
|
|
|
|
|
|
|
|
|
|
class SinusoidalPosEmb(nn.Module):
|
|
|
|
|
def __init__(self, dim):
|
|
|
|
|
super().__init__()
|
|
|
|
|
self.dim = dim
|
|
|
|
|
|
|
|
|
|
def forward(self, x):
|
|
|
|
|
device = x.device
|
|
|
|
|
half_dim = self.dim // 2
|
|
|
|
|
emb = math.log(10000) / (half_dim - 1)
|
|
|
|
|
emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
|
|
|
|
|
emb = x[:, None] * emb[None, :]
|
|
|
|
|
emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
|
|
|
|
|
return emb
|
|
|
|
|
|
|
|
|
|
def Upsample(dim):
|
|
|
|
|
return nn.ConvTranspose2d(dim, dim, 4, 2, 1)
|
|
|
|
|
|
|
|
|
@@ -88,39 +115,6 @@ class PreNorm(nn.Module):
|
|
|
|
|
x = self.norm(x)
|
|
|
|
|
return self.fn(x)
|
|
|
|
|
|
|
|
|
|
# sinusoidal positional embeds
|
|
|
|
|
|
|
|
|
|
class SinusoidalPosEmb(nn.Module):
|
|
|
|
|
def __init__(self, dim):
|
|
|
|
|
super().__init__()
|
|
|
|
|
self.dim = dim
|
|
|
|
|
|
|
|
|
|
def forward(self, x):
|
|
|
|
|
device = x.device
|
|
|
|
|
half_dim = self.dim // 2
|
|
|
|
|
emb = math.log(10000) / (half_dim - 1)
|
|
|
|
|
emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
|
|
|
|
|
emb = x[:, None] * emb[None, :]
|
|
|
|
|
emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
|
|
|
|
|
return emb
|
|
|
|
|
|
|
|
|
|
class LearnedSinusoidalPosEmb(nn.Module):
|
|
|
|
|
""" following @crowsonkb 's lead with learned sinusoidal pos emb """
|
|
|
|
|
""" https://github.com/crowsonkb/v-diffusion-jax/blob/master/diffusion/models/danbooru_128.py#L8 """
|
|
|
|
|
|
|
|
|
|
def __init__(self, dim):
|
|
|
|
|
super().__init__()
|
|
|
|
|
assert (dim % 2) == 0
|
|
|
|
|
half_dim = dim // 2
|
|
|
|
|
self.weights = nn.Parameter(torch.randn(half_dim))
|
|
|
|
|
|
|
|
|
|
def forward(self, x):
|
|
|
|
|
x = rearrange(x, 'b -> b 1')
|
|
|
|
|
freqs = x * rearrange(self.weights, 'd -> 1 d') * 2 * math.pi
|
|
|
|
|
fouriered = torch.cat((freqs.sin(), freqs.cos()), dim = -1)
|
|
|
|
|
fouriered = torch.cat((x, fouriered), dim = -1)
|
|
|
|
|
return fouriered
|
|
|
|
|
|
|
|
|
|
# building block modules
|
|
|
|
|
|
|
|
|
|
class Block(nn.Module):
|
|
|
|
@@ -164,7 +158,6 @@ class ResnetBlock(nn.Module):
|
|
|
|
|
h = self.block1(x, scale_shift = scale_shift)
|
|
|
|
|
|
|
|
|
|
h = self.block2(h)
|
|
|
|
|
|
|
|
|
|
return h + self.res_conv(x)
|
|
|
|
|
|
|
|
|
|
class LinearAttention(nn.Module):
|
|
|
|
@@ -220,6 +213,18 @@ class Attention(nn.Module):
|
|
|
|
|
|
|
|
|
|
# model
|
|
|
|
|
|
|
|
|
|
def MLP(dim_in, dim_hidden):
|
|
|
|
|
return nn.Sequential(
|
|
|
|
|
Rearrange('... -> ... 1'),
|
|
|
|
|
nn.Linear(1, dim_hidden),
|
|
|
|
|
nn.GELU(),
|
|
|
|
|
nn.LayerNorm(dim_hidden),
|
|
|
|
|
nn.Linear(dim_hidden, dim_hidden),
|
|
|
|
|
nn.GELU(),
|
|
|
|
|
nn.LayerNorm(dim_hidden),
|
|
|
|
|
nn.Linear(dim_hidden, dim_hidden)
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
class Unet(nn.Module):
|
|
|
|
|
def __init__(
|
|
|
|
|
self,
|
|
|
|
@@ -230,8 +235,7 @@ class Unet(nn.Module):
|
|
|
|
|
channels = 3,
|
|
|
|
|
resnet_block_groups = 8,
|
|
|
|
|
learned_variance = False,
|
|
|
|
|
learned_sinusoidal_cond = False,
|
|
|
|
|
learned_sinusoidal_dim = 16
|
|
|
|
|
sinusoidal_cond_mlp = True
|
|
|
|
|
):
|
|
|
|
|
super().__init__()
|
|
|
|
|
|
|
|
|
@@ -239,7 +243,7 @@ class Unet(nn.Module):
|
|
|
|
|
|
|
|
|
|
self.channels = channels
|
|
|
|
|
|
|
|
|
|
init_dim = default(init_dim, dim)
|
|
|
|
|
init_dim = default(init_dim, dim // 3 * 2)
|
|
|
|
|
self.init_conv = nn.Conv2d(channels, init_dim, 7, padding = 3)
|
|
|
|
|
|
|
|
|
|
dims = [init_dim, *map(lambda m: dim * m, dim_mults)]
|
|
|
|
@@ -251,21 +255,17 @@ class Unet(nn.Module):
|
|
|
|
|
|
|
|
|
|
time_dim = dim * 4
|
|
|
|
|
|
|
|
|
|
self.learned_sinusoidal_cond = learned_sinusoidal_cond
|
|
|
|
|
self.sinusoidal_cond_mlp = sinusoidal_cond_mlp
|
|
|
|
|
|
|
|
|
|
if learned_sinusoidal_cond:
|
|
|
|
|
sinu_pos_emb = LearnedSinusoidalPosEmb(learned_sinusoidal_dim)
|
|
|
|
|
fourier_dim = learned_sinusoidal_dim + 1
|
|
|
|
|
if sinusoidal_cond_mlp:
|
|
|
|
|
self.time_mlp = nn.Sequential(
|
|
|
|
|
SinusoidalPosEmb(dim),
|
|
|
|
|
nn.Linear(dim, time_dim),
|
|
|
|
|
nn.GELU(),
|
|
|
|
|
nn.Linear(time_dim, time_dim)
|
|
|
|
|
)
|
|
|
|
|
else:
|
|
|
|
|
sinu_pos_emb = SinusoidalPosEmb(dim)
|
|
|
|
|
fourier_dim = dim
|
|
|
|
|
|
|
|
|
|
self.time_mlp = nn.Sequential(
|
|
|
|
|
sinu_pos_emb,
|
|
|
|
|
nn.Linear(fourier_dim, time_dim),
|
|
|
|
|
nn.GELU(),
|
|
|
|
|
nn.Linear(time_dim, time_dim)
|
|
|
|
|
)
|
|
|
|
|
self.time_mlp = MLP(1, time_dim)
|
|
|
|
|
|
|
|
|
|
# layers
|
|
|
|
|
|
|
|
|
@@ -288,8 +288,8 @@ class Unet(nn.Module):
|
|
|
|
|
self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim)))
|
|
|
|
|
self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim = time_dim)
|
|
|
|
|
|
|
|
|
|
for ind, (dim_in, dim_out) in enumerate(reversed(in_out)):
|
|
|
|
|
is_last = ind == (len(in_out) - 1)
|
|
|
|
|
for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):
|
|
|
|
|
is_last = ind >= (num_resolutions - 1)
|
|
|
|
|
|
|
|
|
|
self.ups.append(nn.ModuleList([
|
|
|
|
|
block_klass(dim_out * 2, dim_in, time_emb_dim = time_dim),
|
|
|
|
@@ -301,13 +301,13 @@ class Unet(nn.Module):
|
|
|
|
|
default_out_dim = channels * (1 if not learned_variance else 2)
|
|
|
|
|
self.out_dim = default(out_dim, default_out_dim)
|
|
|
|
|
|
|
|
|
|
self.final_res_block = block_klass(dim * 2, dim, time_emb_dim = time_dim)
|
|
|
|
|
self.final_conv = nn.Conv2d(dim, self.out_dim, 1)
|
|
|
|
|
self.final_conv = nn.Sequential(
|
|
|
|
|
block_klass(dim * 2, dim),
|
|
|
|
|
nn.Conv2d(dim, self.out_dim, 1)
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
def forward(self, x, time):
|
|
|
|
|
x = self.init_conv(x)
|
|
|
|
|
r = x.clone()
|
|
|
|
|
|
|
|
|
|
t = self.time_mlp(time)
|
|
|
|
|
|
|
|
|
|
h = []
|
|
|
|
@@ -330,9 +330,7 @@ class Unet(nn.Module):
|
|
|
|
|
x = attn(x)
|
|
|
|
|
x = upsample(x)
|
|
|
|
|
|
|
|
|
|
x = torch.cat((x, r), dim = 1)
|
|
|
|
|
|
|
|
|
|
x = self.final_res_block(x, t)
|
|
|
|
|
x = torch.cat((x, h.pop()), dim = 1)
|
|
|
|
|
return self.final_conv(x)
|
|
|
|
|
|
|
|
|
|
# gaussian diffusion trainer class
|
|
|
|
@@ -584,22 +582,23 @@ class Trainer(object):
|
|
|
|
|
folder,
|
|
|
|
|
*,
|
|
|
|
|
ema_decay = 0.995,
|
|
|
|
|
image_size = 128,
|
|
|
|
|
train_batch_size = 32,
|
|
|
|
|
train_lr = 1e-4,
|
|
|
|
|
train_num_steps = 100000,
|
|
|
|
|
gradient_accumulate_every = 2,
|
|
|
|
|
amp = False,
|
|
|
|
|
step_start_ema = 2000,
|
|
|
|
|
ema_update_every = 10,
|
|
|
|
|
update_ema_every = 10,
|
|
|
|
|
save_and_sample_every = 1000,
|
|
|
|
|
results_folder = './results',
|
|
|
|
|
augment_horizontal_flip = True
|
|
|
|
|
):
|
|
|
|
|
super().__init__()
|
|
|
|
|
self.image_size = diffusion_model.image_size
|
|
|
|
|
|
|
|
|
|
self.model = diffusion_model
|
|
|
|
|
self.ema = EMA(diffusion_model, beta = ema_decay, update_every = ema_update_every)
|
|
|
|
|
self.ema = EMA(ema_decay)
|
|
|
|
|
self.ema_model = copy.deepcopy(self.model)
|
|
|
|
|
self.update_ema_every = update_ema_every
|
|
|
|
|
|
|
|
|
|
self.step_start_ema = step_start_ema
|
|
|
|
|
self.save_and_sample_every = save_and_sample_every
|
|
|
|
@@ -609,9 +608,9 @@ class Trainer(object):
|
|
|
|
|
self.gradient_accumulate_every = gradient_accumulate_every
|
|
|
|
|
self.train_num_steps = train_num_steps
|
|
|
|
|
|
|
|
|
|
self.ds = Dataset(folder, self.image_size, augment_horizontal_flip = augment_horizontal_flip)
|
|
|
|
|
self.ds = Dataset(folder, image_size, augment_horizontal_flip = augment_horizontal_flip)
|
|
|
|
|
self.dl = cycle(data.DataLoader(self.ds, batch_size = train_batch_size, shuffle = True, pin_memory = True, num_workers = cpu_count()))
|
|
|
|
|
self.opt = Adam(diffusion_model.parameters(), lr = train_lr)
|
|
|
|
|
self.opt = Adam(diffusion_model.parameters(), lr=train_lr)
|
|
|
|
|
|
|
|
|
|
self.step = 0
|
|
|
|
|
|
|
|
|
@@ -621,11 +620,22 @@ class Trainer(object):
|
|
|
|
|
self.results_folder = Path(results_folder)
|
|
|
|
|
self.results_folder.mkdir(exist_ok = True)
|
|
|
|
|
|
|
|
|
|
self.reset_parameters()
|
|
|
|
|
|
|
|
|
|
def reset_parameters(self):
|
|
|
|
|
self.ema_model.load_state_dict(self.model.state_dict())
|
|
|
|
|
|
|
|
|
|
def step_ema(self):
|
|
|
|
|
if self.step < self.step_start_ema:
|
|
|
|
|
self.reset_parameters()
|
|
|
|
|
return
|
|
|
|
|
self.ema.update_model_average(self.ema_model, self.model)
|
|
|
|
|
|
|
|
|
|
def save(self, milestone):
|
|
|
|
|
data = {
|
|
|
|
|
'step': self.step,
|
|
|
|
|
'model': self.model.state_dict(),
|
|
|
|
|
'ema': self.ema.state_dict(),
|
|
|
|
|
'ema': self.ema_model.state_dict(),
|
|
|
|
|
'scaler': self.scaler.state_dict()
|
|
|
|
|
}
|
|
|
|
|
torch.save(data, str(self.results_folder / f'model-{milestone}.pt'))
|
|
|
|
@@ -635,7 +645,7 @@ class Trainer(object):
|
|
|
|
|
|
|
|
|
|
self.step = data['step']
|
|
|
|
|
self.model.load_state_dict(data['model'])
|
|
|
|
|
self.ema.load_state_dict(data['ema'])
|
|
|
|
|
self.ema_model.load_state_dict(data['ema'])
|
|
|
|
|
self.scaler.load_state_dict(data['scaler'])
|
|
|
|
|
|
|
|
|
|
def train(self):
|
|
|
|
@@ -655,15 +665,15 @@ class Trainer(object):
|
|
|
|
|
self.scaler.update()
|
|
|
|
|
self.opt.zero_grad()
|
|
|
|
|
|
|
|
|
|
self.ema.update()
|
|
|
|
|
if self.step % self.update_ema_every == 0:
|
|
|
|
|
self.step_ema()
|
|
|
|
|
|
|
|
|
|
if self.step != 0 and self.step % self.save_and_sample_every == 0:
|
|
|
|
|
self.ema.ema_model.eval()
|
|
|
|
|
with torch.no_grad():
|
|
|
|
|
milestone = self.step // self.save_and_sample_every
|
|
|
|
|
batches = num_to_groups(36, self.batch_size)
|
|
|
|
|
all_images_list = list(map(lambda n: self.ema.ema_model.sample(batch_size=n), batches))
|
|
|
|
|
self.ema_model.eval()
|
|
|
|
|
|
|
|
|
|
milestone = self.step // self.save_and_sample_every
|
|
|
|
|
batches = num_to_groups(36, self.batch_size)
|
|
|
|
|
all_images_list = list(map(lambda n: self.ema_model.sample(batch_size=n), batches))
|
|
|
|
|
all_images = torch.cat(all_images_list, dim=0)
|
|
|
|
|
utils.save_image(all_images, str(self.results_folder / f'sample-{milestone}.png'), nrow = 6)
|
|
|
|
|
self.save(milestone)
|
|
|
|
|
Phil Wang