Changed to blocks API

README.md

@@ -4,7 +4,7 @@ emoji: 🌖
 colorFrom: blue
 colorTo: gray
 sdk: gradio
-sdk_version: 2.9.4
+sdk_version: 3.0b2
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -35,7 +35,21 @@ model_repos = {"e4e": ("akhaliq/JoJoGAN_e4e_ffhq_encode", "e4e_ffhq_encode.pt"),
                "dlib": ("akhaliq/jojogan_dlib", "shape_predictor_68_face_landmarks.dat"),
                "base": ("akhaliq/jojogan-stylegan2-ffhq-config-f", "stylegan2-ffhq-config-f.pt"),
                "anime": ("rinong/stylegan-nada-models", "anime.pt"),
-               "joker": ("rinong/stylegan-nada-models", "joker.pt")
+               "joker": ("rinong/stylegan-nada-models", "joker.pt"),
+               "simpson": ("rinong/stylegan-nada-models", "simpson.pt"),
+               "ssj": ("rinong/stylegan-nada-models", "ssj.pt"),
+               "white_walker": ("rinong/stylegan-nada-models", "white_walker.pt"),
+               "zuckerberg": ("rinong/stylegan-nada-models", "zuckerberg.pt"),
+               "cubism": ("rinong/stylegan-nada-models", "cubism.pt"),
+               "disney_princess": ("rinong/stylegan-nada-models", "disney_princess.pt"),
+               "edvard_munch": ("rinong/stylegan-nada-models", "edvard_munch.pt"),
+               "van_gogh": ("rinong/stylegan-nada-models", "van_gogh.pt"),
+               "oil": ("rinong/stylegan-nada-models", "oil.pt"),
+               "rick_morty": ("rinong/stylegan-nada-models", "rick_morty.pt"),
+               "botero": ("rinong/stylegan-nada-models", "botero.pt"),
+               "crochet": ("rinong/stylegan-nada-models", "crochet.pt"),
+               "modigliani": ("rinong/stylegan-nada-models", "modigliani.pt"),
+               "shrek": ("rinong/stylegan-nada-models", "shrek.pt"),
                }
 
 def get_models():
@@ -113,22 +127,13 @@ class ImageEditor(object):
     def predict(
         self,
         input,              # Input image path
-        output_style,       # Which output style do you want to use?
-        style_list,         # Comma seperated list of models to use. Only accepts models from the output_style list
+        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
     ):  
 
-        if output_style == 'all':
-            styles = self.model_list
-        elif output_style == 'list - enter below':
-            styles = style_list.split(",")
-            for style in styles:
-                if style not in self.model_list:
-                    raise ValueError(f"Encountered style '{style}' in the style_list which is not an available option.")
-        else:
-            styles = [output_style]
+        styles = output_styles
 
         # @title Align image
         input_image = self.run_alignment(str(input))
@@ -200,16 +205,46 @@ class ImageEditor(object):
 
 editor = ImageEditor()
 
-title = "StyleGAN-NADA"
-description = "Gradio Demo for StyleGAN-NADA: CLIP-Guided Domain Adaptation of Image Generators (SIGGRAPH 2022). To use it, upload your image and select a target style. More information about the paper and training new models can be found below."
+def change_component_visibility(visible, components, assign_flipped = False):
+    if assign_flipped:
+        return [component.update(visible=not visible) for component in components]
+    return [component.update(visible=visible) for component in components]
+
+blocks = gr.Blocks()
+
+with blocks:
+    gr.Markdown("<h1><center>StyleGAN-NADA</center></h1>")
+    gr.Markdown(
+        "Demo for StyleGAN-NADA: CLIP-Guided Domain Adaptation of Image Generators (SIGGRAPH 2022)."
+    )
+    gr.Markdown(
+        "For more information about the paper and code for training your own models (with examples OR text), see below."
+    )
+
+    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!")
+        
+        with gr.Row():
+            video_choice = gr.inputs.Checkbox(default=False, label="Generate Video?", optional=False)
+            
+            edit_choice       = gr.inputs.Checkbox(default=False, label="With Editing?", optional=False, visible=False)
+            vid_format_choice = gr.inputs.Radio(choices=["gif", "mp4"], type="value", default='mp4', label="Video Format", visible=False)
+        with gr.Row():
+            img_button = gr.Button("Edit Image")
+            vid_button = gr.Button("Edit Image")
+
+    with gr.Column():
+        img_output = gr.outputs.Image(type="file")
+        vid_output = gr.outputs.Video(visible=False)
+
+    video_choice.change(fn=change_component_visibility, inputs=[video_choice, [edit_choice, vid_format_choice, vid_output, vid_button]], outputs=[edit_choice, vid_format_choice, vid_output, vid_button])
+    video_choice.change(fn=change_component_visibility, inputs=[video_choice, [img_output, img_button], True], outputs=[img_output, img_button])
 
-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>"
+    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)
 
-gr.Interface(editor.predict, [gr.inputs.Image(type="filepath"), 
-                              gr.inputs.Dropdown(choices=editor.get_style_list(), type="value", default='base', label="Model"), 
-                              gr.inputs.Textbox(lines=1, placeholder=None, default="joker,anime,modigliani", label="Style List", optional=True),
-                              gr.inputs.Checkbox(default=False, label="Generate Video?", optional=False),
-                              gr.inputs.Checkbox(default=False, label="With Editing?", optional=False),
-                              gr.inputs.Radio(choices=["gif", "mp4"], type="value", default='mp4', label="Video Format")],
-                              gr.outputs.Image(type="file"), title=title, description=description, article=article, allow_flagging=False, allow_screenshot=False).launch(enable_queue=True)
+    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)
 
+blocks.launch(enable_queue=True)