|
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 project_code_by_edit_name |
|
|
|
import clip |
|
import urllib.request |
|
|
|
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"), |
|
"sketch": ("rinong/stylegan-nada-models", "sketch.pt"), |
|
"santa": ("mjdolan/stylegan-nada-models", "santa.pt"), |
|
"jesus": ("mjdolan/stylegan-nada-models", "jesus.pt"), |
|
"mariah": ("mjdolan/stylegan-nada-models", "mariah.pt"), |
|
"heat_miser": ("mjdolan/stylegan-nada-models", "heat.pt"), |
|
"claymation": ("mjdolan/stylegan-nada-models", "claymation.pt"), |
|
"elf": ("mjdolan/stylegan-nada-models", "elf.pt"), |
|
"krampus": ("mjdolan/stylegan-nada-models", "krampus.pt"), |
|
"grinch": ("mjdolan/stylegan-nada-models", "grinch.pt"), |
|
"jack_frost": ("mjdolan/stylegan-nada-models", "jack_frost.pt"), |
|
"rudolph": ("mjdolan/stylegan-nada-models", "rudolph.pt"), |
|
"home_alone": ("mjdolan/stylegan-nada-models", "home_alone.pt"), |
|
"puppet":("rinong/stylegan-nada-models", "plastic_puppet.pt"), |
|
"crochet": ("rinong/stylegan-nada-models", "crochet.pt"), |
|
"shrek": ("rinong/stylegan-nada-models", "shrek.pt"), |
|
"pixar": ("rinong/stylegan-nada-models", "pixar.pt") |
|
} |
|
|
|
interface_gan_map = {"None": None, "Masculine": ("gender", 1.0), "Feminine": ("gender", -1.0), |
|
"Smiling": ("smile", 1.0), |
|
"Frowning": ("smile", -1.0), "Young": ("age", -1.0), "Old": ("age", 1.0), |
|
"Long Hair": ("hair_length", -1.0), "Short Hair": ("hair_length", 1.0)} |
|
|
|
|
|
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="cuda:0" if torch.cuda.is_available() else "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"] |
|
) |
|
|
|
|
|
self.clip_model, _ = clip.load("ViT-B/32", device=self.device) |
|
|
|
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: |
|
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 get_target_latent(self, source_latent, alter, generators): |
|
np_source_latent = source_latent.squeeze(0).cpu().detach().numpy() |
|
if alter == "None": |
|
return random.choice([source_latent.squeeze(0),] * max((len(generators) - 1), 1)) |
|
edit = interface_gan_map[alter] |
|
projected_code_np = project_code_by_edit_name(np_source_latent, edit[0], edit[1]) |
|
return torch.from_numpy(projected_code_np).float().to(self.device) |
|
|
|
def edit_image(self, input, output_styles, edit_choices): |
|
return self.predict(input, output_styles, edit_choices=edit_choices) |
|
|
|
def predict( |
|
self, |
|
input, |
|
output_styles, |
|
loop_styles=False, |
|
edit_choices=None, |
|
): |
|
|
|
if edit_choices is None: |
|
edit_choices = {"edit_type": "None"} |
|
|
|
|
|
out_dir = tempfile.mkdtemp() |
|
|
|
inverted_latent = self.invert_image(input) |
|
generators = self.get_generators_for_styles(output_styles, loop_styles) |
|
output_paths = [] |
|
|
|
with torch.no_grad(): |
|
for g_ema in generators: |
|
latent_for_gen = self.get_target_latent(inverted_latent, edit_choices, generators) |
|
|
|
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 |
|
|
|
|
|
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() |
|
|
|
img_url = "http://claireye.com.tw/img/230212a.jpg" |
|
urllib.request.urlretrieve(img_url, "pose.jpg") |
|
blocks = gr.Blocks(theme="darkdefault") |
|
|
|
with blocks: |
|
gr.Markdown("<h1><center>Holiday Filters (StyleGAN-NADA)</center></h1>") |
|
gr.Markdown( |
|
"<div>Upload an image of your face, pick your desired output styles, pick any modifiers, and apply StyleGAN-based editing.</div>" |
|
) |
|
with gr.Row(): |
|
with gr.Column(): |
|
input_img = gr.Image(type="filepath", label="Input image") |
|
with gr.Column(): |
|
style_choice = gr.CheckboxGroup(choices=editor.get_style_list(), value=editor.get_style_list(), type="value", label="Styles") |
|
alter = gr.Dropdown( |
|
choices=["None", "Masculine", "Feminine", "Smiling", "Frowning", "Young", "Old", "Short Hair", |
|
"Long Hair"], value="None", label="Additional Modifiers") |
|
img_button = gr.Button("Edit Image") |
|
|
|
with gr.Row(): |
|
img_output = gr.Gallery(label="Output Images") |
|
img_output.style(grid=(3, 3, 4, 4, 6, 6)) |
|
|
|
img_button.click(fn=editor.edit_image, inputs=[input_img, style_choice, alter], outputs=img_output) |
|
ex = gr.Examples(examples=[['pose.jpg', editor.get_style_list(), "Smiling"], ['pose.jpg', editor.get_style_list(), "Long Hair"]], fn=editor.edit_image, inputs=[input_img, style_choice, alter], |
|
outputs=[img_output], cache_examples=True, |
|
run_on_click=True) |
|
ex.dataset.headers = [""] |
|
article = "<p style='text-align: center'><a href='http://claireye.com.tw'>Claireye</a> | 2023</p>" |
|
gr.Markdown(article) |
|
|
|
blocks.launch(enable_queue=True) |
|
|
