Videos now loop.

app.py

@@ -144,22 +144,37 @@ class ImageEditor(object):
 
         return inverted_latent
     
-    def get_generators_for_styles(self, output_styles):
+    def get_generators_for_styles(self, output_styles, loop_styles=False):
+
+        # if style_string:
+        #     styles = style_string.split(",")
+        #     for style in styles:
+        #         if style not in self.model_list:
+        #             raise ValueError(f"Encountered style '{style}' in the input style list which is not an available option.")
+        # else:
+        #     styles = style_checkbox_list
+
+        if "base" in output_styles:                                    # always start with base if chosen
+            output_styles.insert(0, output_styles.pop(output_styles.index("base")))
+        if loop_styles:
+            output_styles.append(output_styles[0])
+
         return [self.generators[style] for style in output_styles]
 
     def edit_image(self, input, output_styles):
-        return self.predict(input, output_styles, False, False, None)
+        return self.predict(input, output_styles)
 
-    def edit_video(self, input, output_styles, with_editing, video_format):
-        return self.predict(input, output_styles, True, with_editing, video_format)
+    def edit_video(self, input, output_styles, with_editing, video_format, loop_styles):
+        return self.predict(input, output_styles, True, with_editing, video_format, loop_styles)
 
     def predict(
         self,
-        input,              # Input image path
-        output_styles,       # Which output style do you want to use?
-        generate_video,     # Generate a video instead of an output image
-        with_editing,       # Apply latent space editing to the generated video
-        video_format        # Choose gif to display in browser, mp4 for higher-quality downloadable video
+        input,                  # Input image path
+        output_styles,          # Style checkbox options.
+        generate_video = False, # Generate a video instead of an output image
+        with_editing   = False, # Apply latent space editing to the generated video
+        video_format   = "mp4", # Choose gif to display in browser, mp4 for higher-quality downloadable video
+        loop_styles    = False, # Loop back to the initial style
     ):  
 
         # @title Align image
@@ -167,7 +182,7 @@ class ImageEditor(object):
         out_path = out_dir / "out.jpg"
         
         inverted_latent = self.invert_image(input)
-        generators = self.get_generators_for_styles(output_styles)
+        generators = self.get_generators_for_styles(output_styles, loop_styles)
 
         if not generate_video:
             with torch.no_grad():
@@ -290,14 +305,15 @@ with blocks:
             with gr.Column():
                 with gr.Row():
                     vid_button = gr.Button("Generate Video")
-                    edit_choice       = gr.inputs.Checkbox(default=False, label="With latent space editing?", optional=False)
+                    loop_styles       = gr.inputs.Checkbox(default=True, label="Loop video back to the initial style?")
+                    edit_choice       = gr.inputs.Checkbox(default=False, label="With latent space editing?")
                     vid_format_choice = gr.inputs.Radio(choices=["gif", "mp4"], type="value", default='mp4', label="Video Format")
             
             with gr.Column():
                 vid_output = gr.outputs.Video(label="Output Video")
 
     img_button.click(fn=editor.edit_image, inputs=[input_img, style_choice], outputs=img_output)
-    vid_button.click(fn=editor.edit_video, inputs=[input_img, style_choice, edit_choice, vid_format_choice], outputs=vid_output)
+    vid_button.click(fn=editor.edit_video, inputs=[input_img, style_choice, edit_choice, vid_format_choice, loop_styles], outputs=vid_output)
 
     article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2108.00946' target='_blank'>StyleGAN-NADA: CLIP-Guided Domain Adaptation of Image Generators</a> | <a href='https://stylegan-nada.github.io/' target='_blank'>Project Page</a> | <a href='https://github.com/rinongal/StyleGAN-nada' target='_blank'>Code</a></p> <center><img src='https://visitor-badge.glitch.me/badge?page_id=rinong_sgnada' alt='visitor badge'></center>"
     gr.Markdown(article)

generate_videos.py

@@ -151,7 +151,7 @@ def video_from_interpolations(fps, output_dir):
                "-r", f"{fps}", 
                "-i", f"{output_dir}/%03d.jpg", 
                "-c:v", "libx264", 
-               "-vf", f"fps={fps}", 
+               "-vf", f'"fps={fps}, loop=loop=-1"', 
                "-pix_fmt", "yuv420p", 
                f"{output_dir}/out.mp4"]