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 from styleclip.styleclip_global import project_code_with_styleclip, style_tensor_to_style_dict 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"), "sc_fs3": ("rinong/stylegan-nada-models", "fs3.npy"), "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") } 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), "Short Hair": ("hair_length", -1.0), "Long 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", "sc_fs3"]] 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.styleclip_fs3 = torch.from_numpy(np.load(model_paths["sc_fs3"])).to(self.device) 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 _pack_edits(func): def inner(self, alter, *args): return func(self, *args, alter) return inner 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) @_pack_edits 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("

Holiday Filters (StyleGAN-NADA)

") gr.Markdown( "
Upload an image of your face, pick your desired output styles, pick any modifiers, and apply StyleGAN-based editing.
" ) 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(), default=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_button.click(fn=editor.edit_image, inputs=[alter, input_img, style_choice], outputs=img_output) article = "

StyleGAN-NADA: CLIP-Guided Domain Adaptation of Image Generators | Project Page | Code

visitor badge
" gr.Markdown(article) blocks.launch(enable_queue=True)