complete video generation code

Author: lucidrains

README.md

@@ -32,9 +32,16 @@ vae = VQGanVAE(
 
 imgs = torch.randn(10, 3, 256, 256)
 
+# alternate learning for autoencoder ...
+
 loss = vae(imgs, return_loss = True)
 loss.backward()
 
+# and the discriminator ...
+
+discr_loss = vae(imgs, return_discr_loss = True)
+discr_loss.backward()
+
 # do above for many steps
 ```
 
@@ -44,22 +51,29 @@ Then, with your learned VAE
 import torch
 from nuwa_pytorch import NUWA, VQGanVAE
 
+# autoencoder
+
 vae = VQGanVAE(
     dim = 512,
-    num_layers = 4
+    num_layers = 4,
+    image_size = 256
 )
 
+# NUWA transformer
+
 nuwa = NUWA(
     vae = vae,
     dim = 512,
     max_video_frames = 5,
     text_num_tokens = 20000,
     image_size = 256
-)
+).cuda()
+
+# data
 
-text = torch.randint(0, 20000, (1, 256))
-mask = torch.ones(1, 256).bool()
-video = torch.randn(1, 5, 3, 256, 256)
+text = torch.randint(0, 20000, (1, 256)).cuda()
+mask = torch.ones(1, 256).bool().cuda()
+video = torch.randn(1, 5, 3, 256, 256).cuda()
 
 loss = nuwa(
     text = text,
@@ -71,17 +85,22 @@ loss = nuwa(
 loss.backward()
 
 # do above with as much data as possible
+
+# then you can generate a video from text
+
+video = nuwa.generate(text = text, text_mask = mask) # (1, 5, 3, 256, 256)
+
 ```
 
 ## Todo
 
 - [x] complete 3dna causal attention in decoder
+- [x] write up easy generation functions
 - [ ] flesh out VAE resnet blocks, offer some choices
 - [ ] make sure GAN portion of VQGan is correct, reread paper
 - [ ] offer new vqvae improvements (orthogonal reg and smaller codebook dimensions)
 - [ ] offer vqvae training script
 - [ ] take care of audio transformer and cross modality attention
-- [ ] write up easy generation functions
 - [ ] segmentation mask encoder, make sure embeddings can undergo 3dna attention with decoder during cross attention
 - [ ] investigate custom attention layouts in microsoft deepspeed sparse attention (using triton)
 

nuwa_pytorch/nuwa_pytorch.py

@@ -24,6 +24,17 @@ def exists(val):
 def default(val, d):
     return val if exists(val) else d
 
+# decorators
+
+def eval_decorator(fn):
+    def inner(model, *args, **kwargs):
+        was_training = model.training
+        model.eval()
+        out = fn(model, *args, **kwargs)
+        model.train(was_training)
+        return out
+    return inner
+
 # tensor helper functions
 
 def log(t, eps = 1e-20):
@@ -129,6 +140,10 @@ def __init__(
         self.discr_loss = hinge_discr_loss if use_hinge_loss else bce_discr_loss
         self.gen_loss = hinge_gen_loss if use_hinge_loss else bce_gen_loss
 
+    @property
+    def codebook(self):
+        return self.vq.codebook
+
     def encode(self, fmap):
         for enc in self.encoders:
             fmap = enc(fmap)
@@ -161,9 +176,11 @@ def forward(
 
         fmap = self.decode(fmap)
 
-        if not return_loss:
+        if not return_loss and not return_discr_loss:
             return fmap
 
+        assert return_loss ^ return_discr_loss, 'you should either return autoencoder loss or discriminator loss, but not both'
+
         if return_discr_loss:
             fmap.detach_()
             fmap_discr_logits, img_discr_logits = map(self.discr, (fmap, img))
@@ -332,16 +349,31 @@ def forward(self, x, mask = None):
         # more variables
 
         kernel_size = self.kernel_size
-        num_frames, fmap_size, _ = self.video_shape
+        fmap_size = self.video_shape[1]
+
+        bos_only = n == 1
+        tokens_per_frame = fmap_size ** 2
+
+        padding = 0 if bos_only else (tokens_per_frame - (n - 1) % tokens_per_frame)
+        num_frames = (n + padding) // tokens_per_frame
 
         # pad for last token in video
 
-        x = F.pad(x, (0, 0, 0, 1), value = 0.)
+        if padding > 0:
+            x = F.pad(x, (0, 0, 0, padding), value = 0.)
 
         # derive queries / keys / values
 
-        qkv = self.to_qkv(x).chunk(3, dim = -1)
-        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h = h), qkv)
+        q, k, v = self.to_qkv(x).chunk(3, dim = -1)
+
+        # early return if <bos>
+
+        if bos_only:
+            return self.to_out(v)
+
+        # split out heads
+
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h = h), (q, k, v))
 
         # scale queries
 
@@ -352,11 +384,6 @@ def forward(self, x, mask = None):
         q = q[:, 1:]
         bos_value = v[:, :1]
 
-        # prepare precomputed causal mask
-
-        causal_mask = self.causal_mask[:n]
-        causal_mask = repeat(causal_mask, 'i j -> b i j', b = b * h)
-
         # compute keys and values
 
         (k_bos, k), (v_bos, v) = map(lambda t: (t[:, :1], t[:, 1:]), (k, v))
@@ -376,6 +403,10 @@ def forward(self, x, mask = None):
 
         # causal mask
 
+        i, j = sim.shape[-2:]
+        causal_mask = self.causal_mask[:i, :j]
+        causal_mask = repeat(causal_mask, 'i j -> b i j', b = b * h)
+
         sim = sim.masked_fill(causal_mask, -torch.finfo(sim.dtype).max)
 
         # attention
@@ -450,6 +481,16 @@ def forward(
 
         return self.norm(x)
 
+# sampling helpers
+
+def top_k(logits, thres = 0.5):
+    num_logits = logits.shape[-1]
+    k = max(int((1 - thres) * num_logits), 1)
+    val, ind = torch.topk(logits, k)
+    probs = torch.full_like(logits, float('-inf'))
+    probs.scatter_(1, ind, val)
+    return probs
+
 # main class
 
 class NUWA(nn.Module):
@@ -495,9 +536,13 @@ def __init__(
 
         fmap_size = image_size // (2 ** vae_num_layers)
 
+        self.video_fmap_size = fmap_size
+        self.max_video_frames = max_video_frames
+        video_shape = (max_video_frames, fmap_size, fmap_size)
+
         self.video_pos_emb = AxialPositionalEmbedding(
             dim = dim,
-            axial_shape = (max_video_frames, fmap_size, fmap_size)
+            axial_shape = video_shape
         )
 
         self.video_transformer = Transformer(
@@ -511,11 +556,66 @@ def __init__(
             ff_dropout = ff_dropout,
             sparse_3dna_attn = True,
             sparse_3dna_kernel_size = sparse_3dna_kernel_size,
-            sparse_3dna_video_shape = (max_video_frames, fmap_size, fmap_size)
+            sparse_3dna_video_shape = video_shape
         )
 
         self.to_logits = nn.Linear(dim, num_image_tokens)
 
+    @torch.no_grad()
+    @eval_decorator
+    def generate(
+        self,
+        *,
+        text,
+        text_mask = None,
+        filter_thres = 0.9,
+        temperature = 1.
+    ):
+        batch, seq_len, device = *text.shape, text.device
+        assert seq_len <= self.text_max_seq_len, 'your input text has a greater length than what was designated on initialization'
+
+        tokens = self.text_embedding(text)
+        pos_emb = self.text_pos_embedding(torch.arange(seq_len, device = device))
+        tokens = tokens + rearrange(pos_emb, 'n d -> 1 n d')
+
+        text_embeds = self.text_transformer(
+            tokens,
+            mask = text_mask
+        )
+
+        bos = repeat(self.video_bos, 'd -> b 1 d', b = batch)
+
+        video_indices = torch.empty((batch, 0), device = device, dtype = torch.long)
+        total_video_tokens = self.video_fmap_size * self.video_fmap_size * self.max_video_frames
+
+        for _ in range(total_video_tokens):
+            frame_embeddings = self.image_embedding(video_indices)
+            frame_embeddings = self.video_pos_emb(frame_embeddings) + frame_embeddings
+            frame_embeddings = torch.cat((bos, frame_embeddings), dim = 1)
+
+            frame_embeddings = self.video_transformer(
+                frame_embeddings,
+                context = text_embeds,
+                context_mask = text_mask
+            )
+
+            logits = self.to_logits(frame_embeddings)
+            logits = logits[:, -1, :]
+
+            filtered_logits = top_k(logits, thres = filter_thres)
+            filtered_logits /= temperature
+            filtered_logits -=  torch.amax(filtered_logits, dim = 1, keepdim = True)
+            probs = F.softmax(filtered_logits, dim = -1)
+            sample = torch.multinomial(probs, 1)
+            video_indices = torch.cat((video_indices, sample), dim = 1)
+
+        codes = self.vae.codebook[video_indices]
+        codes = rearrange(codes, 'b (f h w) d -> (b f) d h w', h = self.video_fmap_size, w = self.video_fmap_size)
+
+        image_reconstructions = self.vae.decode(codes)
+        video = rearrange(image_reconstructions, '(b f) d h w -> b f d h w', b = batch)
+        return video
+
     def forward(
         self,
         *,