Changed blocks format to tabs

app.py

@@ -128,18 +128,7 @@ class ImageEditor(object):
 
         return style_list
 
-    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
-    ):  
-
-        styles = output_styles
-
-        # @title Align image
+    def invert_image(self, input_image):
         input_image = self.run_alignment(str(input))
         
         input_image = input_image.resize(self.resize_dims)
@@ -152,10 +141,32 @@ class ImageEditor(object):
             result_image, latent = images[0], latents[0]
 
         inverted_latent = latent.unsqueeze(0).unsqueeze(1)
+
+        return inverted_latent
+    
+    def get_generators_for_styles(self, output_styles):
+        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)
+
+    def edit_video(self, input, output_styles, with_editing, video_format):
+        return self.predict(input, output_styles, True, with_editing, video_format)
+
+    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
+    ):  
+        # @title Align image
         out_dir = Path(tempfile.mkdtemp())
         out_path = out_dir / "out.jpg"
         
-        generators = [self.generators[style] for style in styles]
+        inverted_latent = self.invert_image(input)
+        generators = self.get_generators_for_styles(output_styles)
 
         if not generate_video:
             with torch.no_grad():
@@ -232,54 +243,22 @@ with blocks:
         "For more information about the paper and code for training your own models (with examples OR text), see below."
     )
 
-    with gr.Row():
-
-        with gr.Column():
-            input_img    = gr.inputs.Image(type="filepath", label="Input image")
-            style_choice = gr.inputs.CheckboxGroup(choices=editor.get_style_list(), type="value", label="Choose your styles!")
-
-            video_choice = gr.inputs.Checkbox(default=False, label="Generate Video?", optional=False)
-            
-            video_options_group = gr.Group()
-            
-            with video_options_group:
-                edit_choice       = gr.inputs.Checkbox(default=False, label="With Editing?", optional=False)
-                vid_format_choice = gr.inputs.Radio(choices=["gif", "mp4"], type="value", default='mp4', label="Video Format")
-        
-            # img_button = gr.Button("Edit Image")
-            # vid_button = gr.Button("Generate Video")
-            img_button = gr.Button("Edit Image")
-            vid_button = gr.Button("Generate Video")
-
-        with gr.Column():
-            img_output = gr.outputs.Image(type="file")
-            vid_output = gr.outputs.Video()
-
-    # visibility_fn = change_component_visibility(component_types=[gr.Checkbox, gr.Radio, gr.Video, gr.Button, gr.Image, gr.Button],
-    #                                             invert_choices=[False, False, False, False, True, True])
-
-    # video_choice.change(fn=visibility_fn, inputs=video_choice, outputs=[edit_choice, vid_format_choice, vid_output, vid_button, img_output, img_button])
-    video_choice.change(fn=group_visibility, inputs=video_choice, outputs=video_options_group)
-    img_button.click(fn=editor.predict, inputs=[input_img, style_choice, video_choice, edit_choice, vid_format_choice], outputs=img_output)
-    vid_button.click(fn=editor.predict, inputs=[input_img, style_choice, video_choice, edit_choice, vid_format_choice], outputs=vid_output)
-    
-    # input_img = gr.inputs.Image(type="filepath", label="Input image")
-    # style_choice = gr.inputs.CheckboxGroup(choices=editor.get_style_list(), type="value", label="Choose your styles!")
-  
-
+    # with gr.Row():
 
-    # with gr.Tabs():
-    #     with gr.TabItem("Edit Images"):
-    #         with gr.Row():
-    #             with gr.Column():
-                    
-            
+    #     with gr.Column():
+    #         input_img    = gr.inputs.Image(type="filepath", label="Input image")
+    #         style_choice = gr.inputs.CheckboxGroup(choices=editor.get_style_list(), type="value", label="Choose your styles!")
 
     #         video_choice = gr.inputs.Checkbox(default=False, label="Generate Video?", optional=False)
             
-    #         edit_choice       = gr.inputs.Checkbox(default=False, label="With Editing?", optional=False)
-    #         vid_format_choice = gr.inputs.Radio(choices=["gif", "mp4"], type="value", default='mp4', label="Video Format")
+    #         video_options_group = gr.Group()
+            
+    #         with video_options_group:
+    #             edit_choice       = gr.inputs.Checkbox(default=False, label="With Editing?", optional=False)
+    #             vid_format_choice = gr.inputs.Radio(choices=["gif", "mp4"], type="value", default='mp4', label="Video Format")
         
+    #         # img_button = gr.Button("Edit Image")
+    #         # vid_button = gr.Button("Generate Video")
     #         img_button = gr.Button("Edit Image")
     #         vid_button = gr.Button("Generate Video")
 
@@ -287,13 +266,36 @@ with blocks:
     #         img_output = gr.outputs.Image(type="file")
     #         vid_output = gr.outputs.Video()
 
-    # visibility_fn = change_component_visibility(component_types=[gr.Checkbox, gr.Radio, gr.Video, gr.Button, gr.Image, gr.Button],
-    #                                             invert_choices=[False, False, False, False, True, True])
-
-    # video_choice.change(fn=visibility_fn, inputs=video_choice, outputs=[edit_choice, vid_format_choice, vid_output, vid_button, img_output, img_button])
+    # # visibility_fn = change_component_visibility(component_types=[gr.Checkbox, gr.Radio, gr.Video, gr.Button, gr.Image, gr.Button],
+    # #                                             invert_choices=[False, False, False, False, True, True])
 
+    # # video_choice.change(fn=visibility_fn, inputs=video_choice, outputs=[edit_choice, vid_format_choice, vid_output, vid_button, img_output, img_button])
+    # video_choice.change(fn=group_visibility, inputs=video_choice, outputs=video_options_group)
     # img_button.click(fn=editor.predict, inputs=[input_img, style_choice, video_choice, edit_choice, vid_format_choice], outputs=img_output)
     # vid_button.click(fn=editor.predict, inputs=[input_img, style_choice, video_choice, edit_choice, vid_format_choice], outputs=vid_output)
+    
+    input_img = gr.inputs.Image(type="filepath", label="Input image")
+
+    with gr.Row():
+        style_choice = gr.inputs.CheckboxGroup(choices=editor.get_style_list(), type="value", label="Choose your styles!")
+
+    with gr.Tabs():
+        with gr.TabItem("Edit Images"):
+            with gr.Row():
+                    img_button = gr.Button("Edit Image")
+                    img_output = gr.outputs.Image(type="file")
+            
+        with gr.TabItem("Create Video"):
+            with gr.Row():
+                edit_choice       = gr.inputs.Checkbox(default=False, label="With latent space editing?", optional=False)
+                vid_format_choice = gr.inputs.Radio(choices=["gif", "mp4"], type="value", default='mp4', label="Video Format")
+
+            with gr.Row():
+                vid_button = gr.Button("Generate Video")
+                vid_output = gr.outputs.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)
 
     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)