complete model portion

denoising_diffusion_pytorch/__init__.py

@@ -1 +1 @@
-from denoising_diffusion_pytorch.denoising_diffusion_pytorch import DenoisingDiffusion
+from denoising_diffusion_pytorch.denoising_diffusion_pytorch import DenoisingDiffusion, Unet

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -1,11 +1,187 @@
+import math
 import torch
 from torch import nn, einsum
 import torch.nn.functional as F
+
+import numpy as np
 from einops import rearrange
 
+# helper models
+
+class Residual(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    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):
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim) * -emb)
+        emb = x[:, None] * emb[None, :]
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+class ResnetBlock(nn.Module):
+    def __init__(self, dim, out_dim, *, time_emb_dim, groups = 32):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            Mish(),
+            nn.Linear(time_emb_dim, out_dim)
+        )
+
+        self.block1 = nn.Sequential(
+            nn.Conv2d(dim, out_dim, 3, padding=1),
+            nn.GroupNorm(groups, out_dim),
+            Mish()
+        )
+
+        self.block2 = nn.Sequential(
+            nn.Conv2d(out_dim, out_dim, 3, padding=1),
+            nn.GroupNorm(groups, out_dim),
+            Mish()
+        )
+
+        self.res_conv = nn.Conv2d(dim, out_dim, 1) if dim != out_dim else nn.Identity()
+
+    def forward(self, x, time_emb):
+        h = self.block1(x)
+        h += self.mlp(time_emb)[:, :, None, None]
+        h = self.block2(h)
+        return h + self.res_conv(x)
+
+class Mish(nn.Module):
+    def forward(self, x):
+        return x * torch.tanh(F.softplus(x))
+
+class Upsample(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.conv = nn.ConvTranspose2d(dim, dim, 4, 2, 1)
+
+    def forward(self, x):
+        return self.conv(x)
+
+class Downsample(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.conv = nn.Conv2d(dim, dim, 3, 2, 1)
+
+    def forward(self, x):
+        return self.conv(x)
+
+class Rezero(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.g = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x):
+        return x * self.g
+
+class LinearAttention(nn.Module):
+    def __init__(self, dim, heads = 8, dim_head = 32):
+        super().__init__()
+        self.heads = heads
+        hidden_dim = dim_head * heads
+        self.to_qkv = nn.Conv2d(dim, hidden_dim, 1, bias = False)
+        self.to_out = nn.Conv2d(hidden_dim, dim, 1)
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+        qkv = self.to_qkv(x)
+        q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads)
+        q = q.softmax(dim=-2)
+        k = k.softmax(dim=-1)
+        context = torch.einsum('bhdn,bhen->bhde', k, v)
+        out = torch.einsum('bhde,bhdn->bhen', context, q)
+        out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w)
+        return self.to_out(out)
+
+# model
+
+class Unet(nn.Module):
+    def __init__(self, dim, dim_mults=(1, 2, 4, 8), groups = 32):
+        super().__init__()
+        dims = [3, *map(lambda m: dim * m, dim_mults)]
+        in_out = list(zip(dims[:-1], dims[1:]))
+
+        self.time_pos_emb = SinusoidalPosEmb(dim)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, dim * 4),
+            Mish(),
+            nn.Linear(dim * 4, dim)
+        )
+
+        self.downs = nn.ModuleList([])
+        self.ups = nn.ModuleList([])
+        num_resolutions = len(in_out)
+
+        for ind, (dim_in, dim_out) in enumerate(in_out):
+            is_last = ind >= (num_resolutions - 1)
+
+            self.downs.append(nn.ModuleList([
+                ResnetBlock(dim_in, dim_out, time_emb_dim = dim),
+                Residual(Rezero(LinearAttention(dim_out))),
+                Downsample(dim_out) if not is_last else nn.Identity()
+            ]))
+
+        mid_dim = dims[-1]
+        self.mid_block1 = ResnetBlock(mid_dim, mid_dim, time_emb_dim = dim)
+        self.mid_attn = Residual(Rezero(LinearAttention(mid_dim)))
+        self.mid_block2 = ResnetBlock(mid_dim, mid_dim, time_emb_dim = dim)
+
+        for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):
+            is_last = ind >= (num_resolutions - 1)
+
+            self.ups.append(nn.ModuleList([
+                ResnetBlock(dim_out * 2, dim_in, time_emb_dim = dim),
+                Residual(Rezero(LinearAttention(dim_in))),
+                Upsample(dim_in) if not is_last else nn.Identity()
+            ]))
+
+        self.final_conv = nn.Sequential(
+            nn.Conv2d(dim, dim, 3, padding = 1),
+            nn.GroupNorm(groups, dim),
+            Mish(),
+            nn.Conv2d(dim, 3, 1)
+        )
+
+    def forward(self, x, time):
+        t = self.time_pos_emb(time)
+        t = self.mlp(t)
+
+        h = []
+
+        for resnet, attn, downsample in self.downs:
+            x = resnet(x, t)
+            x = attn(x)
+            h.append(x)
+            x = downsample(x)
+
+        x = self.mid_block1(x, t)
+        x = self.mid_attn(x)
+        x = self.mid_block2(x, t)
+
+        for resnet, attn, upsample in self.ups:
+            x = torch.cat((x, h.pop()), dim=1)
+            x = resnet(x, t)
+            x = attn(x)
+            x = upsample(x)
+
+        return self.final_conv(x)
+
+# gaussian diffusion trainer class
+
 class DenoisingDiffusion(nn.Module):
     def __init__(self):
         super().__init__()
 
     def forward(self, x):
-        return x
+        return x

setup.py

@@ -14,8 +14,9 @@ setup(
     'generative models'
   ],
   install_requires=[
-    'torch',
-    'einops'
+    'einops',
+    'numpy',
+    'torch'
   ],
   classifiers=[
     'Development Status :: 4 - Beta',