now accepts timestep conditioning, for ddpm

make_a_video_pytorch/make_a_video.py

@@ -1,6 +1,8 @@
-import torch
+import math
 import functools
 from operator import mul
+
+import torch
 from torch import nn, einsum
 
 from einops import rearrange, pack, unpack
@@ -22,6 +24,24 @@ def divisible_by(numer, denom):
 
 mlist = nn.ModuleList
 
+# for time conditioning
+
+class SinusoidalPosEmb(nn.Module):
+ def __init__(self, dim, theta = 10000):
+ super().__init__()
+ self.theta = theta
+ self.dim = dim
+
+ def forward(self, x):
+ dtype, device = x.dtype, x.device
+ assert dtype == torch.float, 'input to sinusoidal pos emb must be a float type'
+
+ half_dim = self.dim // 2
+ emb = math.log(self.theta) / (half_dim - 1)
+ emb = torch.exp(torch.arange(half_dim, device = device, dtype = dtype) * -emb)
+ emb = rearrange(x, 'i -> i 1') * rearrange(emb, 'j -> 1 j')
+ return torch.cat((emb.sin(), emb.cos()), dim = -1).type(dtype)
+
 # layernorm 3d
 
 class LayerNorm(nn.Module):
@@ -226,17 +246,17 @@ def __init__(
  dim,
  dim_out,
  *,
- time_cond_dim = None,
+ timestep_cond_dim = None,
  groups = 8
  ):
  super().__init__()
 
- self.time_mlp = None
+ self.timestep_mlp = None
 
- if exists(time_cond_dim):
- self.time_mlp = nn.Sequential(
+ if exists(timestep_cond_dim):
+ self.timestep_mlp = nn.Sequential(
  nn.SiLU(),
- nn.Linear(time_cond_dim, dim_out * 2)
+ nn.Linear(timestep_cond_dim, dim_out * 2)
  )
 
  self.block1 = Block(dim, dim_out, groups = groups)
@@ -246,14 +266,17 @@ def __init__(
  def forward(
  self,
  x,
- time_emb = None,
+ timestep_emb = None,
  enable_time = True
  ):
+ assert not (exists(timestep_emb) ^ exists(self.timestep_mlp))
 
  scale_shift = None
- if exists(self.time_mlp) and exists(time_emb):
- time_emb = self.time_mlp(time_emb)
- time_emb = rearrange(time_emb, 'b c -> b c 1 1')
+
+ if exists(self.timestep_mlp) and exists(timestep_emb):
+ time_emb = self.timestep_mlp(timestep_emb)
+ to_einsum_eq = 'b c 1 1 1' if x.ndim == 5 else 'b c 1 1'
+ time_emb = rearrange(time_emb, f'b c -> {to_einsum_eq}')
  scale_shift = time_emb.chunk(2, dim = 1)
 
  h = self.block1(x, scale_shift = scale_shift, enable_time = enable_time)
@@ -367,7 +390,8 @@ def __init__(
  self_attns = (False, False, False, True),
  temporal_compression = (False, True, True, True),
  attn_dim_head = 64,
- attn_heads = 8
+ attn_heads = 8,
+ condition_on_timestep = True
  ):
  super().__init__()
  assert len(dim_mult) == len(self_attns) == len(temporal_compression)
@@ -376,6 +400,20 @@ def __init__(
  dims = [dim, *map(lambda mult: mult * dim, dim_mult)]
  dim_in_out = zip(dims[:-1], dims[1:])
 
+ # timestep conditioning for DDPM, not to be confused with the time dimension of the video
+
+ self.to_timestep_cond = None
+ timestep_cond_dim = (dim * 4) if condition_on_timestep else None
+
+ if condition_on_timestep:
+ self.to_timestep_cond = nn.Sequential(
+ SinusoidalPosEmb(dim),
+ nn.Linear(dim, timestep_cond_dim),
+ nn.SiLU()
+ )
+
+ # layers
+
  self.downs = mlist([])
  self.ups = mlist([])
 
@@ -386,21 +424,21 @@ def __init__(
 
  mid_dim = dims[-1]
 
- self.mid_block1 = ResnetBlock(mid_dim, mid_dim)
+ self.mid_block1 = ResnetBlock(mid_dim, mid_dim, timestep_cond_dim = timestep_cond_dim)
  self.mid_attn = SpatioTemporalAttention(dim = mid_dim)
- self.mid_block2 = ResnetBlock(mid_dim, mid_dim)
+ self.mid_block2 = ResnetBlock(mid_dim, mid_dim, timestep_cond_dim = timestep_cond_dim)
 
  for _, self_attend, (dim_in, dim_out), compress_time in zip(range(num_layers), self_attns, dim_in_out, temporal_compression):
 
  self.downs.append(mlist([
- ResnetBlock(dim_in, dim_out),
+ ResnetBlock(dim_in, dim_out, timestep_cond_dim = timestep_cond_dim),
  ResnetBlock(dim_out, dim_out),
  SpatioTemporalAttention(dim = dim_out, **attn_kwargs) if self_attend else None,
  Downsample(dim_out, downsample_time = compress_time)
  ]))
 
  self.ups.append(mlist([
- ResnetBlock(dim_out * 2, dim_in),
+ ResnetBlock(dim_out * 2, dim_in, timestep_cond_dim = timestep_cond_dim),
  ResnetBlock(dim_in, dim_in),
  SpatioTemporalAttention(dim = dim_in, **attn_kwargs) if self_attend else None,
  Upsample(dim_out, upsample_time = compress_time)
@@ -415,14 +453,19 @@ def __init__(
  def forward(
  self,
  x,
+ timestep = None,
  enable_time = True
  ):
+ assert not (exists(self.to_timestep_cond) ^ exists(timestep))
+
+ t = self.to_timestep_cond(rearrange(timestep, '... -> (...)')) if exists(timestep) else None
+
  x = self.conv_in(x, enable_time = enable_time)
 
  hiddens = []
 
  for block1, block2, maybe_attention, downsample in self.downs:
- x = block1(x, enable_time = enable_time)
+ x = block1(x, t, enable_time = enable_time)
  x = block2(x, enable_time = enable_time)
 
  if exists(maybe_attention):
@@ -432,15 +475,15 @@ def forward(
 
  x = downsample(x, enable_time = enable_time)
 
- x = self.mid_block1(x, enable_time = enable_time)
+ x = self.mid_block1(x, t, enable_time = enable_time)
  x = self.mid_attn(x, enable_time = enable_time)
- x = self.mid_block2(x, enable_time = enable_time)
+ x = self.mid_block2(x, t, enable_time = enable_time)
 
  for block1, block2, maybe_attention, upsample in reversed(self.ups):
  x = upsample(x, enable_time = enable_time)
  x = torch.cat((hiddens.pop() * self.skip_scale, x), dim = 1)
 
- x = block1(x, enable_time = enable_time)
+ x = block1(x, t, enable_time = enable_time)
  x = block2(x, enable_time = enable_time)
 
  if exists(maybe_attention):

setup.py

@@ -3,7 +3,7 @@
 setup(
  name = 'make-a-video-pytorch',
  packages = find_packages(exclude=[]),
- version = '0.0.4',
+ version = '0.0.5',
  license='MIT',
  description = 'Make-A-Video - Pytorch',
  author = 'Phil Wang',