add weight standardization prior to groupnorm, lessen cosine sim attention scale to 10 for fp16

README.md

@@ -8,12 +8,8 @@ This implementation was transcribed from the official Tensorflow version <a href
 
 Youtube AI Educators - <a href="https://www.youtube.com/watch?v=W-O7AZNzbzQ">Yannic Kilcher</a> | <a href="https://www.youtube.com/watch?v=344w5h24-h8">AI Coffeebreak with Letitia</a> | <a href="https://www.youtube.com/watch?v=HoKDTa5jHvg">Outlier</a>
 
-<a href="https://github.com/yiyixuxu/denoising-diffusion-flax">Flax implementation</a> from <a href="https://github.com/yiyixuxu">YiYi Xu</a>
-
 <a href="https://huggingface.co/blog/annotated-diffusion">Annotated code</a> by Research Scientists / Engineers from <a href="https://huggingface.co/">🤗 Huggingface</a>
 
-Update: Turns out none of the technicalities really matters at all | <a href="https://arxiv.org/abs/2208.09392">"Cold Diffusion" paper</a>
-
 <img src="./images/sample.png" width="500px"><img>
 
 [![PyPI version](https://badge.fury.io/py/denoising-diffusion-pytorch.svg)](https://badge.fury.io/py/denoising-diffusion-pytorch)
@@ -195,13 +191,3 @@ $ accelerate launch train.py
     volume  = {abs/1903.10520}
 }
 ```
-
-```bibtex
-@article{Salimans2022ProgressiveDF,
-    title   = {Progressive Distillation for Fast Sampling of Diffusion Models},
-    author  = {Tim Salimans and Jonathan Ho},
-    journal = {ArXiv},
-    year    = {2022},
-    volume  = {abs/2202.00512}
-}
-```

denoising_diffusion_pytorch/__init__.py

@@ -4,4 +4,3 @@ from denoising_diffusion_pytorch.learned_gaussian_diffusion import LearnedGaussi
 from denoising_diffusion_pytorch.continuous_time_gaussian_diffusion import ContinuousTimeGaussianDiffusion
 from denoising_diffusion_pytorch.weighted_objective_gaussian_diffusion import WeightedObjectiveGaussianDiffusion
 from denoising_diffusion_pytorch.elucidated_diffusion import ElucidatedDiffusion
-from denoising_diffusion_pytorch.v_param_continuous_time_gaussian_diffusion import VParamContinuousTimeGaussianDiffusion

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -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,11 +53,14 @@ 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
 
+def l2norm(t):
+    return F.normalize(t, dim = -1)
+
 # normalization functions
 
 def normalize_to_neg_one_to_one(img):
@@ -97,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):
@@ -234,12 +239,11 @@ class LinearAttention(nn.Module):
         return self.to_out(out)
 
 class Attention(nn.Module):
-    def __init__(self, dim, heads = 4, dim_head = 32):
+    def __init__(self, dim, heads = 4, dim_head = 32, scale = 10):
         super().__init__()
-        self.scale = dim_head ** -0.5
+        self.scale = scale
         self.heads = heads
         hidden_dim = dim_head * heads
-
         self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)
         self.to_out = nn.Conv2d(hidden_dim, dim, 1)
 
@@ -248,12 +252,11 @@ class Attention(nn.Module):
         qkv = self.to_qkv(x).chunk(3, dim = 1)
         q, k, v = map(lambda t: rearrange(t, 'b (h c) x y -> b h c (x y)', h = self.heads), qkv)
 
-        q = q * self.scale
+        q, k = map(l2norm, (q, k))
 
-        sim = einsum('b h d i, b h d j -> b h i j', q, k)
+        sim = einsum('b h d i, b h d j -> b h i j', q, k) * self.scale
         attn = sim.softmax(dim = -1)
         out = einsum('b h i j, b h d j -> b h i d', attn, v)
-
         out = rearrange(out, 'b h (x y) d -> b (h d) x y', x = h, y = w)
         return self.to_out(out)
 
@@ -429,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
@@ -449,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
@@ -519,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)
 
@@ -573,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 + \
@@ -703,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),
@@ -809,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'])
@@ -844,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()
@@ -854,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()
@@ -871,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')

denoising_diffusion_pytorch/v_param_continuous_time_gaussian_diffusion.py

@@ -1,184 +0,0 @@
-import math
-import torch
-from torch import sqrt
-from torch import nn, einsum
-import torch.nn.functional as F
-from torch.special import expm1
-
-from tqdm import tqdm
-from einops import rearrange, repeat, reduce
-from einops.layers.torch import Rearrange
-
-# helpers
-
-def exists(val):
-    return val is not None
-
-def default(val, d):
-    if exists(val):
-        return val
-    return d() if callable(d) else d
-
-# normalization functions
-
-def normalize_to_neg_one_to_one(img):
-    return img * 2 - 1
-
-def unnormalize_to_zero_to_one(t):
-    return (t + 1) * 0.5
-
-# diffusion helpers
-
-def right_pad_dims_to(x, t):
-    padding_dims = x.ndim - t.ndim
-    if padding_dims <= 0:
-        return t
-    return t.view(*t.shape, *((1,) * padding_dims))
-
-# continuous schedules
-# log(snr) that approximates the original linear schedule
-
-def log(t, eps = 1e-20):
-    return torch.log(t.clamp(min = eps))
-
-def alpha_cosine_log_snr(t, s = 0.008):
-    return -log((torch.cos((t + s) / (1 + s) * math.pi * 0.5) ** -2) - 1, eps = 1e-5)
-
-class VParamContinuousTimeGaussianDiffusion(nn.Module):
-    """
-    a new type of parameterization in v-space proposed in https://arxiv.org/abs/2202.00512 that
-    (1) allows for improved distillation over noise prediction objective and
-    (2) noted in imagen-video to improve upsampling unets by removing the color shifting artifacts
-    """
-
-    def __init__(
-        self,
-        model,
-        *,
-        image_size,
-        channels = 3,
-        num_sample_steps = 500,
-        clip_sample_denoised = True,
-    ):
-        super().__init__()
-        assert model.learned_sinusoidal_cond
-        assert not model.self_condition, 'not supported yet'
-
-        self.model = model
-
-        # image dimensions
-
-        self.channels = channels
-        self.image_size = image_size
-
-        # continuous noise schedule related stuff
-
-        self.log_snr = alpha_cosine_log_snr
-
-        # sampling
-
-        self.num_sample_steps = num_sample_steps
-        self.clip_sample_denoised = clip_sample_denoised        
-
-    @property
-    def device(self):
-        return next(self.model.parameters()).device
-
-    def p_mean_variance(self, x, time, time_next):
-        # reviewer found an error in the equation in the paper (missing sigma)
-        # following - https://openreview.net/forum?id=2LdBqxc1Yv&noteId=rIQgH0zKsRt
-
-        log_snr = self.log_snr(time)
-        log_snr_next = self.log_snr(time_next)
-        c = -expm1(log_snr - log_snr_next)
-
-        squared_alpha, squared_alpha_next = log_snr.sigmoid(), log_snr_next.sigmoid()
-        squared_sigma, squared_sigma_next = (-log_snr).sigmoid(), (-log_snr_next).sigmoid()
-
-        alpha, sigma, alpha_next = map(sqrt, (squared_alpha, squared_sigma, squared_alpha_next))
-
-        batch_log_snr = repeat(log_snr, ' -> b', b = x.shape[0])
-
-        pred_v = self.model(x, batch_log_snr)
-
-        # shown in Appendix D in the paper
-        x_start = alpha * x - sigma * pred_v
-
-        if self.clip_sample_denoised:
-            x_start.clamp_(-1., 1.)
-
-        model_mean = alpha_next * (x * (1 - c) / alpha + c * x_start)
-
-        posterior_variance = squared_sigma_next * c
-
-        return model_mean, posterior_variance
-
-    # sampling related functions
-
-    @torch.no_grad()
-    def p_sample(self, x, time, time_next):
-        batch, *_, device = *x.shape, x.device
-
-        model_mean, model_variance = self.p_mean_variance(x = x, time = time, time_next = time_next)
-
-        if time_next == 0:
-            return model_mean
-
-        noise = torch.randn_like(x)
-        return model_mean + sqrt(model_variance) * noise
-
-    @torch.no_grad()
-    def p_sample_loop(self, shape):
-        batch = shape[0]
-
-        img = torch.randn(shape, device = self.device)
-        steps = torch.linspace(1., 0., self.num_sample_steps + 1, device = self.device)
-
-        for i in tqdm(range(self.num_sample_steps), desc = 'sampling loop time step', total = self.num_sample_steps):
-            times = steps[i]
-            times_next = steps[i + 1]
-            img = self.p_sample(img, times, times_next)
-
-        img.clamp_(-1., 1.)
-        img = unnormalize_to_zero_to_one(img)
-        return img
-
-    @torch.no_grad()
-    def sample(self, batch_size = 16):
-        return self.p_sample_loop((batch_size, self.channels, self.image_size, self.image_size))
-
-    # training related functions - noise prediction
-
-    def q_sample(self, x_start, times, noise = None):
-        noise = default(noise, lambda: torch.randn_like(x_start))
-
-        log_snr = self.log_snr(times)
-
-        log_snr_padded = right_pad_dims_to(x_start, log_snr)
-        alpha, sigma = sqrt(log_snr_padded.sigmoid()), sqrt((-log_snr_padded).sigmoid())
-        x_noised =  x_start * alpha + noise * sigma
-
-        return x_noised, log_snr, alpha, sigma
-
-    def random_times(self, batch_size):
-        return torch.zeros((batch_size,), device = self.device).float().uniform_(0, 1)
-
-    def p_losses(self, x_start, times, noise = None):
-        noise = default(noise, lambda: torch.randn_like(x_start))
-
-        x, log_snr, alpha, sigma = self.q_sample(x_start = x_start, times = times, noise = noise)
-
-        # described in section 4 as the prediction objective, with derivation in Appendix D
-        v = alpha * noise - sigma * x_start
-
-        model_out = self.model(x, log_snr)
-
-        return F.mse_loss(model_out, v)
-
-    def forward(self, img, *args, **kwargs):
-        b, c, h, w, device, img_size, = *img.shape, img.device, self.image_size
-        assert h == img_size and w == img_size, f'height and width of image must be {img_size}'
-
-        times = self.random_times(b)
-        img = normalize_to_neg_one_to_one(img)
-        return self.p_losses(img, times, *args, **kwargs)

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.29.0',
+  version = '0.27.3',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',
@@ -30,4 +30,4 @@ setup(
     'License :: OSI Approved :: MIT License',
     'Programming Language :: Python :: 3.6',
   ],
-)
+)