fp16 version

app.py

@@ -6,9 +6,14 @@ from dalle.models import StoryDalle
 import argparse
 from PIL import Image
 from torchvision.utils import save_image
+import tensorflow as tf
 import tensorflow_hub as hub
 import gdown
+from allennlp.predictors.predictor import Predictor
+import random
 
+torch.set_grad_enabled(False)
+tf.config.set_visible_devices([], 'GPU') # setting Tensorflow's GPU visibility to None to constraing embedding model to CPU
 
 source_frame_paths = {
     'Pororo': '/playpen-ssd/adyasha/projects/StoryGAN/pororo_png/Pororo_ENGLISH1_2/Pororo_ENGLISH1_2_ep6/12.png',
@@ -23,6 +28,51 @@ source_frame_paths = {
 }
 
 
+def get_span_words(span, document):
+    return ' '.join(document[span[0]:span[1]+1])
+
+
+def print_clusters(prediction):
+    document, clusters = prediction['document'], prediction['clusters']
+    for cluster in clusters:
+        print(get_span_words(cluster[0], document) + ': ', end='')
+        print(f"[{'; '.join([get_span_words(span, document) for span in cluster])}]")
+
+
+def resolve_coref(captions, captions_mask, coref_predictor):
+    sent_counts = []
+    doc = ''
+    for cap, mask in zip(captions, captions_mask):
+        if mask == 0:
+            sent_counts.append(0)
+        else:
+            print(cap)
+            count = len([c.strip() for c in cap.split('.') if c.strip()])
+            sent_counts.append(count)
+            doc += cap + ' '
+
+    # print(doc)
+
+    doc = doc.strip()
+    resolved_doc = coref_predictor.coref_resolved(doc)
+    # print(resolved_doc)
+    # print(sent_counts)
+
+    sents = resolved_doc.split('. ')
+    resolved_captions = []
+    for i, (count, mask) in enumerate(zip(sent_counts, captions_mask)):
+        if mask == 0:
+            resolved_captions.append('')
+        else:
+            new_cap = '. '.join(sents[sum(sent_counts[:i]):sum(sent_counts[:i]) + count])
+            new_cap = new_cap.strip()
+            if new_cap[-1] not in ['!', '?', '.']:
+                new_cap += '.'
+            resolved_captions.append(new_cap)
+
+    return resolved_captions
+
+
 def inverse_normalize(tensor, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)):
     mean = torch.as_tensor(mean, dtype=tensor.dtype, device=tensor.device)
     std = torch.as_tensor(std, dtype=tensor.dtype, device=tensor.device)
@@ -66,9 +116,10 @@ def save_story_results(images, video_len=4, n_candidates=1, mask=None):
 
 
 def main(args):
+
     #device = 'cuda:0'
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    #device = torch.device('cpu')
+    # device = torch.device('cpu')
 
     model_url = 'https://drive.google.com/u/1/uc?id=1KAXVtE8lEE2Yc83VY7w6ycOOMkdWbmJo&export=sharing'
 
@@ -81,6 +132,9 @@ def main(args):
     #assert os.path.exists("./ckpt/25.pth")
     gdown.download(png_url, quiet=True, use_cookies=False, output="demo_pororo_good.png")
 
+    coref_model_url = 'https://storage.googleapis.com/allennlp-public-models/coref-spanbert-large-2020.02.27.tar.gz'
+    coref_predictor = Predictor.from_path(coref_model_url)
+
     if args.debug:
         model = None
         embed = None
@@ -88,13 +142,20 @@ def main(args):
         model, config = StoryDalle.from_pretrained(args)
         model.tokenizer.add_tokens(['pororo', 'loopy', 'eddy', 'harry', 'poby', 'tongtong', 'crong', 'rody', 'petty'])
         model.eval()
-        model.to(device=device)
+        # split_model into CPU and GPU
+        if args.split_memory:
+            model.stage2.to(device=device)
+            model.story_linear.to(device=device)
+            model.story_block.to(device=device)
+        else:
+            model.to(device=device)
+            if model.config.story.condition:
+                for i in range(len(model.cross_attention_layers)):
+                    model.cross_attention_layers[i].to(device)
+                print("Cross-attention layers are in cuda:", next(model.cross_attention_layers[0].parameters()).is_cuda)
+
         embed = hub.load("https://tfhub.dev/google/universal-sentence-encoder-large/5")
 
-        if model.config.story.condition:
-            for i in range(len(model.cross_attention_layers)):
-                model.cross_attention_layers[i].to(device)
-            print("Cross-attention layers are in cuda:", next(model.cross_attention_layers[0].parameters()).is_cuda)
 
         valid_transform = transforms.Compose(
             [transforms.Resize(config.dataset.image_resolution),
@@ -103,6 +164,8 @@ def main(args):
              transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]
         )
 
+        print("Model is in ", model.device)
+
     #torch.save(model, './ckpt/checkpoint.pt')
     #sys.exit()
 
@@ -110,32 +173,62 @@ def main(args):
                 supercondition=False):
 
         if not args.debug:
-            captions = [caption_1, caption_2, caption_3, caption_4]
+
+            suffix = random.randint(0, 1000)
+            img_file_path = "./demo/images/gradio_demo_pororo_%s.png" % suffix
+            txt_file_path = "./demo/texts/gradio_demo_pororo_%s.txt" % suffix
+
+            captions = [caption_1.strip(), caption_2.strip(), caption_3.strip(), caption_4.strip()]
+            for i in range(len(captions)):
+                if captions[i][-1] not in ['!', '?', '.']:
+                    captions[i] = captions[i] + '.'
             mask = [1 if caption != '' else 0 for caption in captions]
+
+            with open(txt_file_path, 'w') as f:
+                f.write('\n'.join(captions))
+
             print(captions, mask, source, n_candidates)
+            captions = resolve_coref(captions, mask, coref_predictor)
+            print(captions)
+
             for i, caption in enumerate(captions):
                 if caption == "":
-                    captions[i] = "Pororo is reading a book."
+                    captions[i] = "Pororo is reading a book." # filler for shorter captions
+
             tokens = [model.tokenizer.encode(caption) for caption in captions]
             texts = torch.stack([torch.LongTensor(token.ids) for token in tokens]).unsqueeze(0)
             sent_embeds = torch.tensor(embed(captions).numpy())
-            # sent_embeds = torch.tensor(description_vecs[source_frame_paths[source].
-            #                            replace('/playpen-ssd/adyasha/projects/StoryGAN/pororo_png/', '')[:-4]][0]).unsqueeze(0).repeat(4, 1)
-
             src_image = valid_transform(Image.open('./demo/%s.png' % source).convert('RGB'))
 
             stories = []
             with torch.no_grad():
                 for i in range(texts.shape[0]):
+                    candidates = []
+                    # for _ in range(n_candidates):
+                    #     if args.split_memory: # if splitting model into CPU/GPU, send src_image from CPU memory
+                    #         pixels = model.sampling_batch(texts[i].to(device), src_image.unsqueeze(0),
+                    #                                       sent_embeds.unsqueeze(0).to(device), top_k=top_k, top_p=top_p,
+                    #                                       prompt=None, n_candidates=1, device=device).cpu()
+                    #     else:
+                    #         pixels = model.sampling_batch(texts[i].to(device), src_image.unsqueeze(0).to(device),
+                    #                                       sent_embeds.unsqueeze(0).to(device), top_k=top_k, top_p=top_p,
+                    #                                       prompt=None, n_candidates=1).cpu()
+                    #     print(pixels.shape)
+                    #     candidates.append(pixels.squeeze())
+                    # stories.append(torch.stack(candidates))
+                    #with torch.cuda.amp.autocast():
+
                     pixels = model.sampling_batch(texts[i].to(device), src_image.unsqueeze(0).to(device),
-                                                  sent_embeds.unsqueeze(0).to(device), top_k=top_k, top_p=top_p,
-                                                  prompt=None, n_candidates=n_candidates).cpu()
+                                              sent_embeds.unsqueeze(0).to(device), top_k=top_k, top_p=top_p,
+                                              prompt=None, n_candidates=n_candidates).cpu()
                     stories.append(pixels)
-
             img = save_story_results(stories, video_len=4, n_candidates=n_candidates, mask=mask)
-            save_image(img, "gradio_demo_pororo.png", normalize=True)
+            save_image(img,  img_file_path, normalize=True)
+
+        else:
+            img_file_path = "gradio_demo_pororo.png"
 
-        return "gradio_demo_pororo.png"
+        return img_file_path
 
     with gr.Blocks(css='#output {width:750px; height:750px; float:left;}') as demo:
         gr.Markdown('''
@@ -170,7 +263,7 @@ def main(args):
         Here are some examples of generated visual stories for the above-mentioned settings.
 
         <p align="center">
-            <img src="file/demo_pororo_good.png" width="1000">
+            <img src="file/demo_pororo_good_v1.png" width="1000">
         </p>
 
         Due to the small training dataset size for story visualization, the model has poor generalization to some unseen settings. The model struggles to generate coherent images in the following scenarios.
@@ -236,10 +329,11 @@ def main(args):
         \[4\] Sharma, Piyush, et al. "Conceptual captions: A cleaned, hypernymed, image alt-text dataset for automatic image captioning." Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2018.
         ''')
 
-    demo.launch(share=True)
+    demo.launch(share=False)
 
 
 if __name__ == "__main__":
+
     args_list = ['--model_name_or_path', './ckpt/25.pth',
                  '--prefix_model_name_or_path', './1.3B/',
                  '--dataset_name', 'pororo',
@@ -351,6 +445,7 @@ if __name__ == "__main__":
     )
 
     parser.add_argument("--debug", action="store_true", help="Whether to debug the demo.")
+    parser.add_argument("--split_memory", action="store_true", help="Whether to split the model into GPU & CPU in the demo.")
 
     args = parser.parse_args(args_list)
 

dalle/models/__init__.py

@@ -1094,7 +1094,7 @@ class PromptConditionalDalle(Dalle):
         prompt = self.get_prompt(bsz=5, eval=True)
 
         images = []
-        for t in texts:
+        for i, t in enumerate(texts):
             pixels = self.sampling(t, prompt, top_k=64, num_candidates=5, labels=codes[i]).cpu().numpy()
             pixels = np.transpose(pixels, (0, 2, 3, 1))
             images.append(pixels)
@@ -1211,7 +1211,6 @@ class StoryDalle(Dalle):
                                               lowercase=True,
                                               dropout=None)
 
-
         return model, config_update
 
 
@@ -1224,6 +1223,7 @@ class StoryDalle(Dalle):
                                          resid_pdrop=hparams.resid_pdrop,
                                          attn_pdrop=hparams.attn_pdrop) for i in cross_attention_layers]
 
+
     def get_prompt_p5(self, bsz=None, eval=False):
         input_tokens = self.input_tokens.unsqueeze(0).expand(bsz, -1).to(self.device)
         temp_control = self.wte(input_tokens)
@@ -1232,6 +1232,7 @@ class StoryDalle(Dalle):
             past_key_values = self.dropout(past_key_values)
         return past_key_values
 
+
     def forward(self,
                 images: torch.FloatTensor,
                 src_images: Optional[torch.FloatTensor],
@@ -1287,6 +1288,7 @@ class StoryDalle(Dalle):
         # print(logits_img.shape, logits_txt.shape, codes.shape, texts.shape)
         return logits_img, logits_txt, codes
 
+
     @torch.no_grad()
     def sampling(self,
                  tokens: torch.LongTensor,
@@ -1327,6 +1329,7 @@ class StoryDalle(Dalle):
 
         #with autocast(enabled=False):
         src_codes = self.stage1.get_codes(source).detach()
+
         src_codes = torch.repeat_interleave(src_codes, self.config.story.story_len, dim=0)
         print(tokens.shape, src_codes.shape, prompt.shape)
         if self.config.story.condition:
@@ -1355,6 +1358,7 @@ class StoryDalle(Dalle):
         pixels = torch.clamp(self.stage1.decode_code(codes) * 0.5 + 0.5, 0, 1)  # [B, 256, 256]
         return pixels
 
+
     @torch.no_grad()
     def sampling_batch(self,
                  tokens: torch.LongTensor,
@@ -1363,10 +1367,8 @@ class StoryDalle(Dalle):
                  top_k: int = 256,
                  top_p: Optional[float] = None,
                  softmax_temperature: float = 1.0,
-                 num_candidates: int = 96,
                  device: str = 'cuda:0',
                  use_fp16: bool = True,
-                 labels=None,
                  prompt=None, n_candidates=1) -> torch.FloatTensor:
 
         self.stage1.eval()
@@ -1396,37 +1398,40 @@ class StoryDalle(Dalle):
 
         #with autocast(enabled=False):
         src_codes = self.stage1.get_codes(source).detach()
+        # src_codes = src_codes.to(device=device) #ensure that src_codes is moved to GPU in case VQGAN was kept in CPU
+
+        with torch.autocast(device_type='cuda', dtype=torch.float16):
+            # repeat inputs to adjust to n_candidates and story length
+            src_codes = torch.repeat_interleave(src_codes, self.config.story.story_len * n_candidates, dim=0)
+            prompt = prompt.repeat(n_candidates, 1, 1)
+            pos_enc_prompt = pos_enc_prompt.repeat(n_candidates, 1)
+            tokens = tokens.repeat(n_candidates, 1)
+            print(tokens.shape, src_codes.shape, prompt.shape, pos_enc_prompt.shape)
+            if self.config.story.condition:
+                codes = sampling_conditional(self.stage2,
+                                             self.cross_attention_idxs,
+                                             self.cross_attention_layers,
+                                             tokens,
+                                             src_codes,
+                                             top_k=top_k,
+                                             top_p=top_p,
+                                             softmax_temperature=softmax_temperature,
+                                             use_fp16=use_fp16,
+                                             prompt=prompt,
+                                             pos_prompt=pos_enc_prompt)
+            else:
+                codes = sampling(self.stage2,
+                                 tokens,
+                                 top_k=top_k,
+                                 top_p=top_p,
+                                 softmax_temperature=softmax_temperature,
+                                 use_fp16=use_fp16,
+                                 prompt=prompt,
+                                 pos_prompt=pos_enc_prompt)
+
+            codes = codes.view(self.config.story.story_len * n_candidates, 16, 16)  # [B, 16, 16]
+            print(codes.shape)
 
-        # repeat inputs to adjust to n_candidates and story length
-        src_codes = torch.repeat_interleave(src_codes, self.config.story.story_len * n_candidates, dim=0)
-        prompt = prompt.repeat(n_candidates, 1, 1)
-        pos_enc_prompt = pos_enc_prompt.repeat(n_candidates, 1)
-        tokens = tokens.repeat(n_candidates, 1)
-        print(tokens.shape, src_codes.shape, prompt.shape, pos_enc_prompt.shape)
-        if self.config.story.condition:
-            codes = sampling_conditional(self.stage2,
-                                         self.cross_attention_idxs,
-                                         self.cross_attention_layers,
-                                         tokens,
-                                         src_codes,
-                                         top_k=top_k,
-                                         top_p=top_p,
-                                         softmax_temperature=softmax_temperature,
-                                         use_fp16=use_fp16,
-                                         prompt=prompt,
-                                         pos_prompt=pos_enc_prompt)
-        else:
-            codes = sampling(self.stage2,
-                             tokens,
-                             top_k=top_k,
-                             top_p=top_p,
-                             softmax_temperature=softmax_temperature,
-                             use_fp16=use_fp16,
-                             prompt=prompt,
-                             pos_prompt=pos_enc_prompt)
-
-        codes = codes.view(self.config.story.story_len * n_candidates, 16, 16)  # [B, 16, 16]
-        print(codes.shape)
         pixels = torch.clamp(self.stage1.decode_code(codes) * 0.5 + 0.5, 0, 1)  # [B, 3, 256, 256]
         print(pixels.shape)
         return pixels.view(n_candidates, self.config.story.story_len, pixels.shape[-3], pixels.shape[-2], pixels.shape[-1])
@@ -1444,11 +1449,10 @@ class StoryDalle(Dalle):
         pred = pred.view(bs, 16, 16)  # [B, 16, 16]
         pixels = torch.clamp(self.stage1.decode_code(pred) * 0.5 + 0.5, 0, 1).cpu().numpy()  # [B, 256, 256]
         pixels = np.transpose(pixels, (0, 2, 3, 1))
-
         prompt = self.get_prompt(bsz=5, eval=True)
 
         images = []
-        for t in texts:
+        for i, t in enumerate(texts):
             pixels = self.sampling(t, prompt, top_k=64, num_candidates=5, labels=codes[i]).cpu().numpy()
             pixels = np.transpose(pixels, (0, 2, 3, 1))
             images.append(pixels)

dalle/models/stage2/layers.py

@@ -182,8 +182,13 @@ class Block(nn.Module):
     def sample_with_context(self, x, context, context_mask, cross_attn_layer, layer_past=None):
         attn, present = self.attn(self.ln1(x), use_cache=True, layer_past=layer_past)
         x = x + attn
-        c_attn = cross_attn_layer(x, context, context_mask)
-        x = x + c_attn
+
+        c_attn = cross_attn_layer(x.to(device=context.device),
+                                  context,
+                                  context_mask.to(device=context.device))
+
+        x = x + c_attn.to(device=x.device)
+
         x = x + self.mlp(self.ln2(x))
         return x, present
 

requirements.txt

@@ -10,3 +10,4 @@ pytorch-lightning
 einops
 tokenizers
 tensorflow
+allennlp==2.10.0