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("StyleGAN-NADA
")
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), not 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), not 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 = "StyleGAN-NADA: CLIP-Guided Domain Adaptation of Image Generators | Project Page | Code
"
gr.Markdown(article)
blocks.launch(enable_queue=True)