Overhauled editing UI, output to gallery

app.py

@@ -1,33 +1,24 @@
 import os
-from posixpath import basename
+import random
 
 import torch
 import gradio as gr
 
-import os
-import sys
-import numpy as np
-
 from e4e.models.psp import pSp
 from util import *
 from huggingface_hub import hf_hub_download
 
-import os
-import sys
 import tempfile
-import shutil
 from argparse import Namespace
-from pathlib import Path
 import shutil
 
 import dlib
 import numpy as np
 import torchvision.transforms as transforms
 from torchvision import utils
-from PIL import Image
 
 from model.sg2_model import Generator
-from generate_videos import generate_frames, video_from_interpolations, vid_to_gif
+from generate_videos import generate_frames, video_from_interpolations, project_code_by_edit_name
 
 model_dir = "models"
 os.makedirs(model_dir, exist_ok=True)
@@ -120,7 +111,6 @@ class ImageEditor(object):
         print("setup complete")
 
     def get_style_list(self):
-        # style_list = ['all', 'list - enter below']
         style_list = []
 
         for key in self.generators:
@@ -146,26 +136,70 @@ class ImageEditor(object):
     
     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
+        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)
+    def _pack_edits(func):
+        def inner(self, 
+                  edit_type_choice, 
+                  pose_slider, 
+                  smile_slider, 
+                  gender_slider, 
+                  age_slider, 
+                  hair_slider, 
+                  src_text_styleclip, 
+                  tar_text_styleclip, 
+                  alpha_styleclip, 
+                  beta_styleclip,
+                  *args):
+
+            edit_choices = {"edit_type": edit_type_choice,
+                            "pose": pose_slider,
+                            "smile": smile_slider,
+                            "gender": gender_slider,
+                            "age": age_slider,
+                            "hair": hair_slider,
+                            "src_text": src_text_styleclip,
+                            "tar_text": tar_text_styleclip,
+                            "alpha": alpha_styleclip,
+                            "beta": beta_styleclip}
+
+            return func(self, *args, edit_choices)
+
+        return inner
+
+    def get_target_latents(self, source_latent, edit_choices, generators):
+
+        np_source_latent = source_latent.squeeze(0).cpu().detach().numpy()
+
+        target_latents = []
+
+        if edit_choices["edit_type"] == "InterFaceGAN":
+            for attribute_name in ["pose", "smile", "gender", "age", "hair"]:
+                strength = edit_choices[attribute_name]
+                if strength != 0.0:
+                    target_latents.append(project_code_by_edit_name(np_source_latent, attribute_name, strength))
+
+        elif edit_choices["edit_type"] == "StyleCLIP":
+            pass
+        
+        # if edit type is none or if all slides were set to 0
+        if not target_latents:
+            target_latents = [source_latent, ] * (len(generators) - 1)
+        
+        return target_latents
+
+    @_pack_edits
+    def edit_image(self, input, output_styles, edit_choices):
+        return self.predict(input, output_styles, edit_choices)
 
-    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)
+    @_pack_edits
+    def edit_video(self, input, output_styles, loop_styles, edit_choices):
+        return self.predict(input, output_styles, True, loop_styles, edit_choices)
 
     def predict(
         self,
@@ -173,55 +207,57 @@ class ImageEditor(object):
         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
+        edit_choices   = None,  # Optional dictionary with edit choice arguments
     ):  
 
+        if edit_choices is None:
+            edit_choices = {"edit_type": "None"}
+
         # @title Align image
-        out_dir = Path(tempfile.mkdtemp())
-        out_path = out_dir / "out.jpg"
+        out_dir = tempfile.mkdtemp()
         
         inverted_latent = self.invert_image(input)
         generators = self.get_generators_for_styles(output_styles, loop_styles)
 
+        target_latents = self.get_target_latents(inverted_latent, edit_choices, generators)
+
         if not generate_video:
+            output_paths = []
+
             with torch.no_grad():
-                img_list = []
                 for g_ema in generators:
-                    img, _ = g_ema(inverted_latent, input_is_latent=True, truncation=1, randomize_noise=False)
-                    img_list.append(img)
-                
-                out_img = torch.cat(img_list, axis=0)
-                utils.save_image(out_img, out_path, nrow=int(np.sqrt(out_img.size(0))), normalize=True, scale_each=True, range=(-1, 1))
+                    latent_for_gen = random.choice(target_latents)
+                    latent_for_gen = [torch.from_numpy(latent_for_gen).float().to(self.device)]
 
-            return str(out_path)
+                    img, _ = g_ema(latent_for_gen, input_is_latent=True, truncation=1, randomize_noise=False)
+
+                    output_path = os.path.join(out_dir, f"out_{len(output_paths)}.jpg")
+                    utils.save_image(img, output_path, nrow=1, normalize=True, range=(-1, 1))
+
+                    output_paths.append(output_path)
+
+            return output_paths
         
-        return self.generate_vid(generators, inverted_latent, out_dir, video_format, with_editing)
-
-    def generate_vid(self, generators, latent, out_dir, video_format, with_editing):      
-        np_latent = latent.squeeze(0).cpu().detach().numpy()
-        args = {
-                'fps': 24,
-                'target_latents': None,
-                'edit_directions': None,
-                'unedited_frames': 0 if with_editing else 40 * (len(generators) - 1)
-                }
-
-        args = Namespace(**args)
+        return self.generate_vid(generators, inverted_latent, out_dir, with_editing)
+
+    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(args, np_latent, generators, dirpath)
-            video_from_interpolations(args.fps, dirpath)
+            generate_frames(np_latent, target_latents, generators, dirpath)
+            video_from_interpolations(fps, dirpath)
             
-            gen_path = Path(dirpath) / "out.mp4"
-            out_path = out_dir / f"out.{video_format}"
+            gen_path = os.path.join(dirpath, "out.mp4")
+            out_path = os.path.join(out_dir, "out.mp4")
 
-            if video_format == 'gif':
-                vid_to_gif(gen_path, out_dir, scale=256, fps=args.fps)
-            else:
-                shutil.copy2(gen_path, out_path) 
+            shutil.copy2(gen_path, out_path) 
 
-        return str(out_path)
+        return out_path
 
     def run_alignment(self, image_path):
         aligned_image = align_face(filepath=image_path, predictor=self.shape_predictor)
@@ -236,12 +272,12 @@ class ImageEditor(object):
 
 editor = ImageEditor()
 
-def change_component_visibility(component_types, invert_choices):
+# def change_component_visibility(component_types, invert_choices):
 
-    def visibility_impl(visible):
-        return [component_types[idx].update(visible=visible ^ invert_choices[idx]) for idx in range(len(component_types))]
+#     def visibility_impl(visible):
+#         return [component_types[idx].update(visible=visible ^ invert_choices[idx]) for idx in range(len(component_types))]
     
-    return visibility_impl
+#     return visibility_impl
 
 # def group_visibility(visible):
 #     print("visible: ", visible)
@@ -258,60 +294,59 @@ with blocks:
     gr.Markdown(
         "For more information about the paper and code for training your own models (with examples OR text), see below."
     )
-
+    
     with gr.Row():
+        input_img = gr.inputs.Image(type="filepath", label="Input image")
 
-        with gr.Column():
-            input_img    = gr.inputs.Image(type="filepath", label="Input image")
+        with gr.Column():    
             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)
-            
-            loop_styles       = gr.inputs.Checkbox(default=True, label="Loop video back to the initial style?", visible=False)
-            edit_choice       = gr.inputs.Checkbox(default=False, label="With Editing?", visible=False)
-            vid_format_choice = gr.inputs.Radio(choices=["gif", "mp4"], type="value", default='mp4', label="Video Format", visible=False)
-        
-            # img_button = gr.Button("Edit Image")
-            # vid_button = gr.Button("Generate Video")
-            img_button = gr.Button("Edit Image")
-            vid_button = gr.Button("Generate Video", visible=False)
-
-        with gr.Column():
-            img_output = gr.outputs.Image(type="file")
-            vid_output = gr.outputs.Video(visible=False)
-
-    visibility_fn = change_component_visibility(component_types=[gr.Checkbox, gr.Radio, gr.Video, gr.Button, gr.Image, gr.Button, gr.Checkbox],
-                                                invert_choices=[False, False, False, False, True, True, False])
-
-    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.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, loop_styles], outputs=vid_output)
-    
-    # with gr.Row():
-    #     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.Tabs():
-    #     with gr.TabItem("Edit Images"):
-    #         with gr.Column():
-    #             img_button = gr.Button("Edit Image")
-    #         with gr.Column():
-    #             img_output = gr.outputs.Image(type="file", label="Output Image")
+            editing_type_choice = gr.Radio(choices=["None", "InterFaceGAN", "StyleCLIP"], label="Choose latent space editing option. For InterFaceGAN and StyleCLIP, set the options below:")
             
-    #     with gr.TabItem("Create Video"):
-    #         with gr.Column():
-    #             with gr.Row():
-    #                 vid_button = gr.Button("Generate Video")
-    #                 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")
+            with gr.Tabs():
+                with gr.TabItem("InterFaceGAN Editing Options"):
+                    gr.Markdown("Move the sliders to make the chosen attribute stronger (e.g. the person older) or leave at 0 to disable editing.")
+                    gr.Markdown("If multiple options are provided, they will be used randomly between images (or sequentially for a video), <u>not</u> together")
+
+                    pose_slider   = gr.Slider(label="Pose", minimum=-1, maximum=1, value=0, step=0.02)
+                    smile_slider  = gr.Slider(label="Smile", minimum=-1, maximum=1, value=0, step=0.02)
+                    gender_slider = gr.Slider(label="Perceived Gender", minimum=-1, maximum=1, value=0, step=0.02)
+                    age_slider    = gr.Slider(label="Age", minimum=-1, maximum=1, value=0, step=0.02)
+                    hair_slider   = gr.Slider(label="Hair Length", minimum=-1, maximum=1, value=0, step=0.02)
+
+                    ig_edit_choices = [pose_slider, smile_slider, gender_slider, age_slider, hair_slider]
+
+                with gr.TabItem("StyleCLIP Editing Options"):
+                    gr.Markdown("Move the sliders to make the chosen attribute stronger (e.g. the person older) or leave at 0 to disable editing.")
+                    gr.Markdown("If multiple options are provided, they will be used randomly between images (or sequentially for a video), <u>not</u> together")
 
-    # 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, loop_styles], outputs=vid_output)
+                    src_text_styleclip = gr.Textbox(label="Source text")
+                    tar_text_styleclip = gr.Textbox(label="Target text")
+
+                    alpha_styleclip    = gr.Slider(label="Edit strength", minimum=-10, maximum=10, value=0, step=0.1)
+                    beta_styleclip     = gr.Slider(label="Disentanglement Threshold", minimum=0.08, maximum=0.3, value=0.14, step=0.01)
+
+                    sc_edit_choices = [src_text_styleclip, tar_text_styleclip, alpha_styleclip, beta_styleclip]
+
+    with gr.Tabs():
+        with gr.TabItem("Edit Images"):
+            with gr.Column():
+                img_button = gr.Button("Edit Image")
+            with gr.Column():
+                img_output = gr.Gallery(label="Output Images")
+            
+        with gr.TabItem("Create Video"):
+            with gr.Row():
+                with gr.Column():
+                    vid_button  = gr.Button("Generate Video")
+                    loop_styles = gr.inputs.Checkbox(default=True, label="Loop video back to the initial style?")
+                    
+                with gr.Column():
+                    vid_output = gr.outputs.Video(label="Output Video")
+
+    edit_inputs = [editing_type_choice] + ig_edit_choices + sc_edit_choices
+    img_button.click(fn=editor.edit_image, inputs=edit_inputs + [input_img, style_choice], outputs=img_output)
+    vid_button.click(fn=editor.edit_video, inputs=edit_inputs + [input_img, style_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

@@ -35,12 +35,12 @@ import copy
 VALID_EDITS = ["pose", "age", "smile", "gender", "hair_length", "beard"]
 
 SUGGESTED_DISTANCES = {
-                       "pose": (3.0, -3.0),
-                       "smile": (2.0, -2.0),
-                       "age": (4.0, -4.0),
-                       "gender": (3.0, -3.0),
-                       "hair_length": (None, -4.0),
-                       "beard": (2.0, None)
+                       "pose": 3.0,
+                       "smile": 2.0,
+                       "age": 4.0,
+                       "gender": 3.0,
+                       "hair_length": -4.0,
+                       "beard": 2.0
                       }
                     
 def project_code(latent_code, boundary, distance=3.0):
@@ -50,21 +50,26 @@ def project_code(latent_code, boundary, distance=3.0):
 
     return latent_code + distance * boundary
 
-def generate_frames(args, source_latent, g_ema_list, output_dir):
+def project_code_by_edit_name(latent_code, name, strength):
+    boundary_dir = Path(os.path.abspath(__file__)).parents[0].joinpath("editing", "interfacegan_boundaries")
+
+    distance = SUGGESTED_DISTANCES[name] * strength
+    boundary = torch.load(os.path.join(boundary_dir, f'{name}.pt'), map_location="cpu").numpy()
+
+    return project_code(latent_code, boundary, distance)
+
+def generate_frames(source_latent, target_latents, g_ema_list, output_dir):
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    alphas = np.linspace(0, 1, num=20)
-    
-    interpolate_func = interpolate_with_boundaries # default
-    if args.target_latents:                        # if provided with targets
-        interpolate_func = interpolate_with_target_latents
-    if args.unedited_frames:                       # if only interpolating through generators
-        interpolate_func = duplicate_latent
+    num_alphas = min(20, 60 // len(target_latents))
 
-    latents = interpolate_func(args, source_latent, alphas)
+    alphas = np.linspace(0, 1, num=num_alphas)
+    
+    latents = interpolate_with_target_latents(source_latent, target_latents, alphas)
 
     segments = len(g_ema_list) - 1
+
     if segments:
         segment_length = len(latents) / segments
 
@@ -96,50 +101,15 @@ def generate_frames(args, source_latent, g_ema_list, output_dir):
 def interpolate_forward_backward(source_latent, target_latent, alphas):
     latents_forward  = [a * target_latent + (1-a) * source_latent for a in alphas] # interpolate from source to target
     latents_backward = latents_forward[::-1]                                       # interpolate from target to source
-    return latents_forward + [target_latent] * 20 + latents_backward               # forward + short delay at target + return
+    return latents_forward + [target_latent] * len(alphas) + latents_backward      # forward + short delay at target + return
 
-def duplicate_latent(args, source_latent, alphas):
-    return [source_latent for _ in range(args.unedited_frames)]
-
-def interpolate_with_boundaries(args, source_latent, alphas):
-    edit_directions = args.edit_directions or ['pose', 'smile', 'gender', 'age', 'hair_length']
-    
-    # interpolate latent codes with all targets
-
-    print("Interpolating latent codes...")
-
-    boundary_dir = Path(os.path.abspath(__file__)).parents[0].joinpath("editing", "interfacegan_boundaries")
-
-    boundaries_and_distances = []
-    for direction_type in edit_directions:
-        distances = SUGGESTED_DISTANCES[direction_type]
-        boundary = torch.load(os.path.join(boundary_dir, f'{direction_type}.pt'), map_location="cpu").numpy()
-
-        for distance in distances:
-            if distance:
-                boundaries_and_distances.append((boundary, distance))
-
-    latents = []
-    for boundary, distance in boundaries_and_distances:
-        
-        target_latent = project_code(source_latent, boundary, distance)
-        latents.extend(interpolate_forward_backward(source_latent, target_latent, alphas)) 
-
-    return latents
-
-def interpolate_with_target_latents(args, source_latent, alphas):    
+def interpolate_with_target_latents(source_latent, target_latents, alphas):    
     # interpolate latent codes with all targets
 
     print("Interpolating latent codes...")
     
     latents = []
-    for target_latent_path in args.target_latents:
-        
-        if target_latent_path == args.source_latent:
-            continue
-
-        target_latent = np.load(target_latent_path, allow_pickle=True)
-
+    for target_latent in target_latents:
         latents.extend(interpolate_forward_backward(source_latent, target_latent, alphas))
 
     return latents
@@ -157,105 +127,5 @@ def video_from_interpolations(fps, output_dir):
     
     subprocess.call(command)
 
-def merge_videos(output_dir, num_subdirs):
-
-    output_file = os.path.join(output_dir, "combined.mp4")
-
-    if num_subdirs == 1: # if we only have one video, just copy it over
-        shutil.copy2(os.path.join(output_dir, str(0), "out.mp4"), output_file)
-    else:                # otherwise merge using ffmpeg
-        command = ["ffmpeg"]
-        for dir in range(num_subdirs):
-            command.extend(['-i', os.path.join(output_dir, str(dir), "out.mp4")])
-        
-        sqrt_subdirs = int(num_subdirs ** .5)
-
-        if (sqrt_subdirs ** 2) != num_subdirs:
-            raise ValueError("Number of checkpoints cannot be arranged in a square grid")
-        
-        command.append("-filter_complex")
-
-        filter_string = ""
-        vstack_string = ""
-        for row in range(sqrt_subdirs):
-            row_str = ""
-            for col in range(sqrt_subdirs):
-                row_str += f"[{row * sqrt_subdirs + col}:v]"
-
-            letter = chr(ord('A')+row)
-            row_str += f"hstack=inputs={sqrt_subdirs}[{letter}];"
-            vstack_string += f"[{letter}]"
-
-            filter_string += row_str
-        
-        vstack_string += f"vstack=inputs={sqrt_subdirs}[out]"
-        filter_string += vstack_string
-
-        command.extend([filter_string, "-map", "[out]", output_file])
-
-        subprocess.call(command)
-
-def vid_to_gif(vid_path, output_dir, scale=256, fps=35):
-
-    command = ["ffmpeg", 
-               "-i", f"{vid_path}", 
-               "-vf", f"fps={fps},scale={scale}:-1:flags=lanczos,split[s0][s1];[s0]palettegen[p];[s1]fifo[s2];[s2][p]paletteuse",  
-               "-loop", "0",
-               f"{output_dir}/out.gif"]
-    
-    subprocess.call(command)
-
-
-if __name__ == '__main__':
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument('--size', type=int, default=1024)
-    parser.add_argument('--ckpt', type=str, nargs="+", required=True, help="Path to one or more pre-trained generator checkpoints.")
-    parser.add_argument('--channel_multiplier', type=int, default=2)
-    parser.add_argument('--out_dir', type=str, required=True, help="Directory where output files will be placed")
-    parser.add_argument('--source_latent', type=str, required=True, help="Path to an .npy file containing an initial latent code")
-    parser.add_argument('--target_latents', nargs="+", type=str, help="A list of paths to .npy files containing target latent codes to interpolate towards, or a directory containing such .npy files.")
-    parser.add_argument('--force', '-f', action='store_true', help="Force run with non-empty directory. Image files not overwritten by the proccess may still be included in the final video")
-    parser.add_argument('--fps', default=35, type=int, help='Frames per second in the generated videos.')
-    parser.add_argument('--edit_directions', nargs="+", type=str, help=f"A list of edit directions to use in video generation (if not using a target latent directory). Available directions are: {VALID_EDITS}")
-    parser.add_argument('--unedited_frames', type=int, default=0, help="Used to generate videos with no latent editing. If set to a positive number and target_latents is not provided, will simply duplicate the initial frame <unedited_frames> times.")
-    
-    args = parser.parse_args()
-
-    os.makedirs(args.out_dir, exist_ok=True)
-
-    if not args.force and os.listdir(args.out_dir):
-        print("Output directory is not empty. Either delete the directory content or re-run with -f.")
-        exit(0)
-
-    if args.target_latents and len(args.target_latents) == 1 and os.path.isdir(args.target_latents[0]):
-        args.target_latents = [os.path.join(args.target_latents[0], file_name) for file_name in os.listdir(args.target_latents[0]) if file_name.endswith(".npy")]
-        args.target_latents = sorted(args.target_latents)
-
-    args.latent = 512
-    args.n_mlp = 8
-
-    g_ema = Generator(
-        args.size, args.latent, args.n_mlp, channel_multiplier=args.channel_multiplier
-    ).to(device)
-
-    source_latent = np.load(args.source_latent, allow_pickle=True)
-
-    for idx, ckpt_path in enumerate(args.ckpt):
-        print(f"Generating video using checkpoint: {ckpt_path}")
-        checkpoint = torch.load(ckpt_path)
-
-        g_ema.load_state_dict(checkpoint['g_ema'])
-
-        output_dir = os.path.join(args.out_dir, str(idx))
-        os.makedirs(output_dir)
-
-        generate_frames(args, source_latent, [g_ema], output_dir)
-        video_from_interpolations(args.fps, output_dir)
-
-    merge_videos(args.out_dir, len(args.ckpt))
-