|
|
|
@@ -37,9 +37,6 @@ def default(val, d):
|
|
|
|
|
return val
|
|
|
|
|
return d() if callable(d) else d
|
|
|
|
|
|
|
|
|
|
def identity(t, *args, **kwargs):
|
|
|
|
|
return t
|
|
|
|
|
|
|
|
|
|
def cycle(dl):
|
|
|
|
|
while True:
|
|
|
|
|
for data in dl:
|
|
|
|
@@ -56,7 +53,7 @@ def num_to_groups(num, divisor):
|
|
|
|
|
arr.append(remainder)
|
|
|
|
|
return arr
|
|
|
|
|
|
|
|
|
|
def convert_image_to_fn(img_type, image):
|
|
|
|
|
def convert_image_to(img_type, image):
|
|
|
|
|
if image.mode != img_type:
|
|
|
|
|
return image.convert(img_type)
|
|
|
|
|
return image
|
|
|
|
@@ -100,11 +97,16 @@ class WeightStandardizedConv2d(nn.Conv2d):
|
|
|
|
|
eps = 1e-5 if x.dtype == torch.float32 else 1e-3
|
|
|
|
|
|
|
|
|
|
weight = self.weight
|
|
|
|
|
mean = reduce(weight, 'o ... -> o 1 1 1', 'mean')
|
|
|
|
|
var = reduce(weight, 'o ... -> o 1 1 1', partial(torch.var, unbiased = False))
|
|
|
|
|
normalized_weight = (weight - mean) * (var + eps).rsqrt()
|
|
|
|
|
flattened_weights = rearrange(weight, 'o ... -> o (...)')
|
|
|
|
|
|
|
|
|
|
return F.conv2d(x, normalized_weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
|
|
|
|
|
mean = reduce(weight, 'o ... -> o 1 1 1', 'mean')
|
|
|
|
|
|
|
|
|
|
var = torch.var(flattened_weights, dim = -1, unbiased = False)
|
|
|
|
|
var = rearrange(var, 'o -> o 1 1 1')
|
|
|
|
|
|
|
|
|
|
weight = (weight - mean) * (var + eps).rsqrt()
|
|
|
|
|
|
|
|
|
|
return F.conv2d(x, weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
|
|
|
|
|
|
|
|
|
|
class LayerNorm(nn.Module):
|
|
|
|
|
def __init__(self, dim):
|
|
|
|
@@ -430,7 +432,6 @@ class GaussianDiffusion(nn.Module):
|
|
|
|
|
):
|
|
|
|
|
super().__init__()
|
|
|
|
|
assert not (type(self) == GaussianDiffusion and model.channels != model.out_dim)
|
|
|
|
|
assert not model.learned_sinusoidal_cond
|
|
|
|
|
|
|
|
|
|
self.model = model
|
|
|
|
|
self.channels = self.model.channels
|
|
|
|
@@ -450,7 +451,7 @@ class GaussianDiffusion(nn.Module):
|
|
|
|
|
raise ValueError(f'unknown beta schedule {beta_schedule}')
|
|
|
|
|
|
|
|
|
|
alphas = 1. - betas
|
|
|
|
|
alphas_cumprod = torch.cumprod(alphas, dim=0)
|
|
|
|
|
alphas_cumprod = torch.cumprod(alphas, axis=0)
|
|
|
|
|
alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value = 1.)
|
|
|
|
|
|
|
|
|
|
timesteps, = betas.shape
|
|
|
|
@@ -520,19 +521,16 @@ class GaussianDiffusion(nn.Module):
|
|
|
|
|
posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
|
|
|
|
|
return posterior_mean, posterior_variance, posterior_log_variance_clipped
|
|
|
|
|
|
|
|
|
|
def model_predictions(self, x, t, x_self_cond = None, clip_x_start = False):
|
|
|
|
|
def model_predictions(self, x, t, x_self_cond = None):
|
|
|
|
|
model_output = self.model(x, t, x_self_cond)
|
|
|
|
|
maybe_clip = partial(torch.clamp, min = -1., max = 1.) if clip_x_start else identity
|
|
|
|
|
|
|
|
|
|
if self.objective == 'pred_noise':
|
|
|
|
|
pred_noise = model_output
|
|
|
|
|
x_start = self.predict_start_from_noise(x, t, pred_noise)
|
|
|
|
|
x_start = maybe_clip(x_start)
|
|
|
|
|
x_start = self.predict_start_from_noise(x, t, model_output)
|
|
|
|
|
|
|
|
|
|
elif self.objective == 'pred_x0':
|
|
|
|
|
pred_noise = self.predict_noise_from_start(x, t, model_output)
|
|
|
|
|
x_start = model_output
|
|
|
|
|
x_start = maybe_clip(x_start)
|
|
|
|
|
pred_noise = self.predict_noise_from_start(x, t, x_start)
|
|
|
|
|
|
|
|
|
|
return ModelPrediction(pred_noise, x_start)
|
|
|
|
|
|
|
|
|
@@ -574,30 +572,31 @@ class GaussianDiffusion(nn.Module):
|
|
|
|
|
def ddim_sample(self, shape, clip_denoised = True):
|
|
|
|
|
batch, device, total_timesteps, sampling_timesteps, eta, objective = shape[0], self.betas.device, self.num_timesteps, self.sampling_timesteps, self.ddim_sampling_eta, self.objective
|
|
|
|
|
|
|
|
|
|
times = torch.linspace(-1, total_timesteps - 1, steps=sampling_timesteps + 1) # [-1, 0, 1, 2, ..., T-1] when sampling_timesteps == total_timesteps
|
|
|
|
|
times = torch.linspace(0., total_timesteps, steps = sampling_timesteps + 2)[:-1]
|
|
|
|
|
times = list(reversed(times.int().tolist()))
|
|
|
|
|
time_pairs = list(zip(times[:-1], times[1:])) # [(T-1, T-2), (T-2, T-3), ..., (1, 0), (0, -1)]
|
|
|
|
|
time_pairs = list(zip(times[:-1], times[1:]))
|
|
|
|
|
|
|
|
|
|
img = torch.randn(shape, device = device)
|
|
|
|
|
|
|
|
|
|
x_start = None
|
|
|
|
|
|
|
|
|
|
for time, time_next in tqdm(time_pairs, desc = 'sampling loop time step'):
|
|
|
|
|
time_cond = torch.full((batch,), time, device=device, dtype=torch.long)
|
|
|
|
|
alpha = self.alphas_cumprod_prev[time]
|
|
|
|
|
alpha_next = self.alphas_cumprod_prev[time_next]
|
|
|
|
|
|
|
|
|
|
time_cond = torch.full((batch,), time, device = device, dtype = torch.long)
|
|
|
|
|
|
|
|
|
|
self_cond = x_start if self.self_condition else None
|
|
|
|
|
pred_noise, x_start, *_ = self.model_predictions(img, time_cond, self_cond, clip_x_start = clip_denoised)
|
|
|
|
|
|
|
|
|
|
if time_next < 0:
|
|
|
|
|
img = x_start
|
|
|
|
|
continue
|
|
|
|
|
pred_noise, x_start, *_ = self.model_predictions(img, time_cond, self_cond)
|
|
|
|
|
|
|
|
|
|
alpha = self.alphas_cumprod[time]
|
|
|
|
|
alpha_next = self.alphas_cumprod[time_next]
|
|
|
|
|
if clip_denoised:
|
|
|
|
|
x_start.clamp_(-1., 1.)
|
|
|
|
|
|
|
|
|
|
sigma = eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt()
|
|
|
|
|
c = (1 - alpha_next - sigma ** 2).sqrt()
|
|
|
|
|
c = ((1 - alpha_next) - sigma ** 2).sqrt()
|
|
|
|
|
|
|
|
|
|
noise = torch.randn_like(img)
|
|
|
|
|
noise = torch.randn_like(img) if time_next > 0 else 0.
|
|
|
|
|
|
|
|
|
|
img = x_start * alpha_next.sqrt() + \
|
|
|
|
|
c * pred_noise + \
|
|
|
|
@@ -704,7 +703,7 @@ class Dataset(Dataset):
|
|
|
|
|
self.image_size = image_size
|
|
|
|
|
self.paths = [p for ext in exts for p in Path(f'{folder}').glob(f'**/*.{ext}')]
|
|
|
|
|
|
|
|
|
|
maybe_convert_fn = partial(convert_image_to_fn, convert_image_to) if exists(convert_image_to) else nn.Identity()
|
|
|
|
|
maybe_convert_fn = partial(convert_image_to, convert_image_to) if exists(convert_image_to) else nn.Identity()
|
|
|
|
|
|
|
|
|
|
self.transform = T.Compose([
|
|
|
|
|
T.Lambda(maybe_convert_fn),
|
|
|
|
@@ -810,10 +809,7 @@ class Trainer(object):
|
|
|
|
|
torch.save(data, str(self.results_folder / f'model-{milestone}.pt'))
|
|
|
|
|
|
|
|
|
|
def load(self, milestone):
|
|
|
|
|
accelerator = self.accelerator
|
|
|
|
|
device = accelerator.device
|
|
|
|
|
|
|
|
|
|
data = torch.load(str(self.results_folder / f'model-{milestone}.pt'), map_location=device)
|
|
|
|
|
data = torch.load(str(self.results_folder / f'model-{milestone}.pt'))
|
|
|
|
|
|
|
|
|
|
model = self.accelerator.unwrap_model(self.model)
|
|
|
|
|
model.load_state_dict(data['model'])
|
|
|
|
@@ -845,7 +841,6 @@ class Trainer(object):
|
|
|
|
|
|
|
|
|
|
self.accelerator.backward(loss)
|
|
|
|
|
|
|
|
|
|
accelerator.clip_grad_norm_(self.model.parameters(), 1.0)
|
|
|
|
|
pbar.set_description(f'loss: {total_loss:.4f}')
|
|
|
|
|
|
|
|
|
|
accelerator.wait_for_everyone()
|
|
|
|
@@ -855,7 +850,6 @@ class Trainer(object):
|
|
|
|
|
|
|
|
|
|
accelerator.wait_for_everyone()
|
|
|
|
|
|
|
|
|
|
self.step += 1
|
|
|
|
|
if accelerator.is_main_process:
|
|
|
|
|
self.ema.to(device)
|
|
|
|
|
self.ema.update()
|
|
|
|
@@ -872,6 +866,7 @@ class Trainer(object):
|
|
|
|
|
utils.save_image(all_images, str(self.results_folder / f'sample-{milestone}.png'), nrow = int(math.sqrt(self.num_samples)))
|
|
|
|
|
self.save(milestone)
|
|
|
|
|
|
|
|
|
|
self.step += 1
|
|
|
|
|
pbar.update(1)
|
|
|
|
|
|
|
|
|
|
accelerator.print('training complete')
|
|
|
|
|
Phil Wang