Merge pull request #110 from lorcalhost/main

Add save GIF animation feature

Author: lucidrains

README.md

@@ -168,6 +168,9 @@ FLAGS
     --theta_hidden=THETA_INITIAL
         Default: 30.0
         Hyperparameter describing the frequency of the color space. Only applies to the hidden layers of the network.
+    --save_gif=SAVE_GIF
+        Default: False
+        Wether or not to save a GIF animation of the generation procedure. Only works if save_progress is set to True.
 ```
 
 ### Priming

deep_daze/cli.py

@@ -33,6 +33,7 @@ def train(
         create_story=False,
         story_start_words=5,
         story_words_per_epoch=5,
+        save_gif=False
 ):
     """
     :param text: (required) A phrase less than 77 characters which you would like to visualize.
@@ -62,6 +63,7 @@ def train(
     :param create_story: Creates a story by optimizing each epoch on a new sliding-window of the input words. If this is enabled, much longer texts than 77 chars can be used. Requires save_progress to visualize the transitions of the story.
     :param story_start_words: Only used if create_story is True. How many words to optimize on for the first epoch.
     :param story_words_per_epoch: Only used if create_story is True. How many words to add to the optimization goal per epoch after the first one.
+    :param save_gif: Only used if save_progress is True. Saves a GIF animation of the generation procedure using the saved frames.
     """
     # Don't instantiate imagine if the user just wants help.
     if any("--help" in arg for arg in sys.argv):
@@ -95,7 +97,8 @@ def train(
         saturate_bound=saturate_bound,
         create_story=create_story,
         story_start_words=story_start_words,
-        story_words_per_epoch=story_words_per_epoch
+        story_words_per_epoch=story_words_per_epoch,
+        save_gif=save_gif
     )
 
     print('Starting up...')

deep_daze/deep_daze.py

@@ -13,6 +13,7 @@
 from torch_optimizer import DiffGrad, AdamP
 
 from PIL import Image
+from imageio import imread, mimsave
 import torchvision.transforms as T
 #from torchvision.utils import save_image
 
@@ -223,6 +224,7 @@ def __init__(
             create_story=False,
             story_start_words=5,
             story_words_per_epoch=5,
+            save_gif=False
     ):
 
         super().__init__()
@@ -293,6 +295,8 @@ def __init__(
 
             image_tensor = self.clip_img_transform(image)[None, ...].cuda()
             self.start_image = image_tensor
+
+        self.save_gif = save_gif
 
     def create_clip_encoding(self, text=None, img=None, encoding=None):
         self.text = text
@@ -410,6 +414,15 @@ def save_image(self, epoch, iteration, img=None):
 
         tqdm.write(f'image updated at "./{str(self.filename)}"')
 
+    def generate_gif(self):
+        images = []
+        for file_name in sorted(os.listdir('./')):
+            if file_name.startswith(self.textpath) and file_name != f'{self.textpath}.jpg':
+                images.append(imread(os.path.join('./', file_name)))
+
+        mimsave(f'{self.textpath}.gif', images)
+        print(f'Generated image generation animation at ./{self.textpath}.gif')
+
     def forward(self):
         if exists(self.start_image):
             tqdm.write('Preparing with initial image...')
@@ -453,3 +466,7 @@ def forward(self):
                 self.clip_encoding = self.update_story_encoding(epoch, i)
 
         self.save_image(epoch, i) # one final save at end
+
+        if self.save_gif and self.save_progress:
+            self.generate_gif()
+        

setup.py

@@ -30,6 +30,7 @@
     'einops>=0.3',
     'fire',
     'ftfy',
+    'imageio>=2.9.0',
     'siren-pytorch>=0.0.8',
     'torch>=1.7.1',
     'torch_optimizer',