app.py

import os
import random

import torch
import gradio as gr

from e4e.models.psp import pSp
from util import *
from huggingface_hub import hf_hub_download

import tempfile
from argparse import Namespace
import shutil

import dlib
import numpy as np
import torchvision.transforms as transforms
from torchvision import utils

from model.sg2_model import Generator
from generate_videos import generate_frames, video_from_interpolations, project_code_by_edit_name

model_dir = "models"
os.makedirs(model_dir, exist_ok=True)

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"),
            #    "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"),
            #    "sketch": ("rinong/stylegan-nada-models", "sketch.pt"),
            #    "thanos": ("rinong/stylegan-nada-models", "thanos.pt"),
               }

def get_models():
    os.makedirs(model_dir, exist_ok=True)

    model_paths = {}

    for model_name, repo_details in model_repos.items():
        download_path = hf_hub_download(repo_id=repo_details[0], filename=repo_details[1])
        model_paths[model_name] = download_path

    return model_paths

model_paths = get_models()

class ImageEditor(object):
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        latent_size = 512
        n_mlp = 8
        channel_mult = 2
        model_size = 1024

        self.generators = {}

        self.model_list = [name for name in model_paths.keys() if name not in ["e4e", "dlib"]]

        for model in self.model_list:
            g_ema = Generator(
                model_size, latent_size, n_mlp, channel_multiplier=channel_mult
            ).to(self.device)

            checkpoint = torch.load(model_paths[model], map_location=self.device)

            g_ema.load_state_dict(checkpoint['g_ema'])

            self.generators[model] = g_ema

        self.experiment_args = {"model_path": model_paths["e4e"]}
        self.experiment_args["transform"] = transforms.Compose(
            [
                transforms.Resize((256, 256)),
                transforms.ToTensor(),
                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
            ]
        )
        self.resize_dims = (256, 256)

        model_path = self.experiment_args["model_path"]

        ckpt = torch.load(model_path, map_location="cpu")
        opts = ckpt["opts"]

        opts["checkpoint_path"] = model_path
        opts = Namespace(**opts)

        self.e4e_net = pSp(opts, self.device)
        self.e4e_net.eval()

        self.shape_predictor = dlib.shape_predictor(
            model_paths["dlib"]
        )

        print("setup complete")

    def get_style_list(self):
        style_list = []

        for key in self.generators:
            style_list.append(key)

        return style_list

    def invert_image(self, input_image):
        input_image = self.run_alignment(str(input_image))
        
        input_image = input_image.resize(self.resize_dims)

        img_transforms = self.experiment_args["transform"]
        transformed_image = img_transforms(input_image)

        with torch.no_grad():
            images, latents = self.run_on_batch(transformed_image.unsqueeze(0))
            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, loop_styles=False):

        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 _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_length": 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_length"]:
                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 = [np_source_latent, ] * max((len(generators) - 1), 1)
        
        return target_latents

    @_pack_edits
    def edit_image(self, input, output_styles, edit_choices):
        return self.predict(input, output_styles, edit_choices=edit_choices)

    @_pack_edits
    def edit_video(self, input, output_styles, loop_styles, edit_choices):
        return self.predict(input, output_styles, generate_video=True, loop_styles=loop_styles, edit_choices=edit_choices)

    def predict(
        self,
        input,                  # Input image path
        output_styles,          # Style checkbox options.
        generate_video = False, # Generate a video instead of an output image
        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 = 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():
                for g_ema in generators:
                    latent_for_gen = random.choice(target_latents)
                    latent_for_gen = [torch.from_numpy(latent_for_gen).float().to(self.device)]

                    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, 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)
            video_from_interpolations(fps, dirpath)
            
            gen_path = os.path.join(dirpath, "out.mp4")
            out_path = os.path.join(out_dir, "out.mp4")

            shutil.copy2(gen_path, out_path) 

        return out_path

    def run_alignment(self, image_path):
        aligned_image = align_face(filepath=image_path, predictor=self.shape_predictor)
        print("Aligned image has shape: {}".format(aligned_image.size))
        return aligned_image

    def run_on_batch(self, inputs):
        images, latents = self.e4e_net(
            inputs.to(self.device).float(), randomize_noise=False, return_latents=True
        )
        return images, latents

editor = ImageEditor()

# 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))]
    
#     return visibility_impl

# def group_visibility(visible):
#     print("visible: ", visible)

#     return gr.Group.update(visibile=visible)

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.Row():
        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!")

            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.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.05)
                    smile_slider  = gr.Slider(label="Smile", minimum=-1, maximum=1, value=0, step=0.05)
                    gender_slider = gr.Slider(label="Perceived Gender", minimum=-1, maximum=1, value=0, step=0.05)
                    age_slider    = gr.Slider(label="Age", minimum=-1, maximum=1, value=0, step=0.05)
                    hair_slider   = gr.Slider(label="Hair Length", minimum=-1, maximum=1, value=0, step=0.05)

                    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")

                    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.Row():
                with gr.Column():
                    with gr.Row():
                        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():
                    with gr.Row():
                        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)

blocks.launch(enable_queue=True)