allow user to set number of interpolation steps

Author: lucidrains

README.md

@@ -76,7 +76,7 @@ $ stylegan2_pytorch  --generate
 To generate a video of a interpolation through two random points in latent space.
 
 ```bash
-$ stylegan2_pytorch --generate-interpolation
+$ stylegan2_pytorch --generate-interpolation --interpolation-num-steps 100
 ```
 
 To save each individual frame of the interpolation

setup.py

@@ -8,7 +8,7 @@
           'stylegan2_pytorch = stylegan2_pytorch.cli:main',
       ],
   },
-  version = '1.0.0',
+  version = '1.0.1',
   license='GPLv3+',
   description = 'StyleGan2 in Pytorch',
   author = 'Phil Wang',

stylegan2_pytorch/cli.py

@@ -82,6 +82,7 @@ def train_from_folder(
     save_every = 1000,
     generate = False,
     generate_interpolation = False,
+    interpolation_num_steps = 100,
     save_frames = False,
     num_image_tiles = 8,
     trunc_psi = 0.75,
@@ -135,7 +136,7 @@ def train_from_folder(
         model = Trainer(**model_args)
         model.load(load_from)
         samples_name = timestamped_filename()
-        model.generate_interpolation(samples_name, num_image_tiles, save_frames = save_frames)
+        model.generate_interpolation(samples_name, num_image_tiles, num_steps = interpolation_num_steps, save_frames = save_frames)
         print(f'interpolation generated at {results_dir}/{name}/{samples_name}')
         return
 

stylegan2_pytorch/stylegan2_pytorch.py

@@ -998,7 +998,7 @@ def generate_truncated(self, S, G, style, noi, trunc_psi = 0.75, num_image_tiles
         return generated_images.clamp_(0., 1.)
 
     @torch.no_grad()
-    def generate_interpolation(self, num = 0, num_image_tiles = 8, trunc = 1.0, save_frames = False):
+    def generate_interpolation(self, num = 0, num_image_tiles = 8, trunc = 1.0, num_steps = 100, save_frames = False):
         self.GAN.eval()
         ext = 'jpg' if not self.transparent else 'png'
         num_rows = num_image_tiles
@@ -1009,11 +1009,11 @@ def generate_interpolation(self, num = 0, num_image_tiles = 8, trunc = 1.0, save
 
         # latents and noise
 
-        latents_low = noise(num_rows ** 2, latent_dim)
-        latents_high = noise(num_rows ** 2, latent_dim)
+        latents_low = noise(num_rows ** 2, latent_dim, device=self.rank)
+        latents_high = noise(num_rows ** 2, latent_dim, device=self.rank)
         n = image_noise(num_rows ** 2, image_size, device=self.rank)
 
-        ratios = torch.linspace(0., 8., 100)
+        ratios = torch.linspace(0., 8., num_steps)
 
         frames = []
         for ratio in tqdm(ratios):