Spaces:
Build error
Build error
| 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 | |
| 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: # 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 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, # Input image path | |
| output_styles, # Style checkbox options. | |
| 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) | |
| 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() | |
| 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=[['example1.jpg', editor.get_style_list(), "Smiling"], ['example2.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='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=mjdolan.holiday_stylegan_nada' alt='visitor badge'></center>" | |
| gr.Markdown(article) | |
| blocks.launch(enable_queue=True) | |