Added styleclip support to video generation

app.py

@@ -258,18 +258,15 @@ class ImageEditor(object):
                     output_paths.append(output_path)
 
             return output_paths
-        
+
         return self.generate_vid(generators, inverted_latent, target_latents, out_dir)
 
     def generate_vid(self, generators, source_latent, target_latents, out_dir):    
 
         fps = 24
 
-        np_latent = source_latent.squeeze(0).cpu().detach().numpy()
-
         with tempfile.TemporaryDirectory() as dirpath:
-
-            generate_frames(np_latent, target_latents, generators, dirpath)
+            generate_frames(source_latent, target_latents, generators, dirpath)
             video_from_interpolations(fps, dirpath)
             
             gen_path = os.path.join(dirpath, "out.mp4")

generate_videos.py

@@ -32,6 +32,8 @@ import subprocess
 import shutil
 import copy
 
+from styleclip.styleclip_global import style_tensor_to_style_dict, style_dict_to_style_tensor
+
 VALID_EDITS = ["pose", "age", "smile", "gender", "hair_length", "beard"]
 
 SUGGESTED_DISTANCES = {
@@ -62,14 +64,24 @@ def generate_frames(source_latent, target_latents, g_ema_list, output_dir):
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
+    code_is_s = target_latents.size()[1] == 9088
+
+    if code_is_s:
+        source_s_dict = g_ema_list[0].get_s_code(source_latent, input_is_latent=True)[0]
+        np_latent = style_dict_to_style_tensor(source_s_dict, g_ema_list[0]).cpu().detach().numpy()
+        target_latents = target_latents.cpu().detach().numpy()
+    else:
+        np_latent = source_latent.squeeze(0).cpu().detach().numpy()
+
+
     num_alphas = min(10, 30 // len(target_latents))
 
     alphas = np.linspace(0, 1, num=num_alphas)
     
-    latents = interpolate_with_target_latents(source_latent, target_latents, alphas)
+    latents = interpolate_with_target_latents(np_latent, target_latents, alphas)
 
     segments = len(g_ema_list) - 1
-    
+
     if segments:
         segment_length = len(latents) / segments
 
@@ -91,10 +103,14 @@ def generate_frames(source_latent, target_latents, g_ema_list, output_dir):
                 src_pars[k].data.copy_(mix_pars[segment_id][k] * (1 - mix_alpha) + mix_pars[segment_id + 1][k] * mix_alpha)
 
         if idx == 0 or segments or latent is not latents[idx - 1]:
-            w = torch.from_numpy(latent).float().to(device)
+            latent_tensor = torch.from_numpy(latent).float().to(device)
 
             with torch.no_grad():
-                img, _ = g_ema([w], input_is_latent=True, truncation=1, randomize_noise=False)
+                if code_is_s:
+                    latent_for_gen = style_tensor_to_style_dict(latent_tensor, g_ema)
+                    img, _ = g_ema(latent_for_gen, input_is_s_code=True, input_is_latent=True, truncation=1, randomize_noise=False)
+                else:
+                    img, _ = g_ema([latent_tensor], input_is_latent=True, truncation=1, randomize_noise=False)
 
         utils.save_image(img, f"{output_dir}/{str(idx).zfill(3)}.jpg", nrow=1, normalize=True, scale_each=True, range=(-1, 1))