|
import os |
|
from PIL import Image |
|
import torch |
|
import gradio as gr |
|
import torch |
|
torch.backends.cudnn.benchmark = True |
|
from torchvision import transforms, utils |
|
from util import * |
|
from PIL import Image |
|
import math |
|
import random |
|
import numpy as np |
|
from torch import nn, autograd, optim |
|
from torch.nn import functional as F |
|
from tqdm import tqdm |
|
import lpips |
|
from model import * |
|
from copy import deepcopy |
|
import imageio |
|
|
|
import os |
|
import sys |
|
import numpy as np |
|
from PIL import Image |
|
import torch |
|
import torchvision.transforms as transforms |
|
from argparse import Namespace |
|
from e4e.models.psp import pSp |
|
from util import * |
|
from huggingface_hub import hf_hub_download |
|
|
|
device= 'cpu' |
|
model_path_e = hf_hub_download(repo_id="akhaliq/JoJoGAN_e4e_ffhq_encode", filename="e4e_ffhq_encode.pt") |
|
ckpt = torch.load(model_path_e, map_location='cpu') |
|
opts = ckpt['opts'] |
|
opts['checkpoint_path'] = model_path_e |
|
opts= Namespace(**opts) |
|
net = pSp(opts, device).eval().to(device) |
|
|
|
@ torch.no_grad() |
|
def projection(img, name, device='cuda'): |
|
|
|
transform = transforms.Compose( |
|
[ |
|
transforms.Resize(256), |
|
transforms.CenterCrop(256), |
|
transforms.ToTensor(), |
|
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]), |
|
] |
|
) |
|
img = transform(img).unsqueeze(0).to(device) |
|
images, w_plus = net(img, randomize_noise=False, return_latents=True) |
|
result_file = {} |
|
result_file['latent'] = w_plus[0] |
|
torch.save(result_file, name) |
|
return w_plus[0] |
|
|
|
device = 'cpu' |
|
|
|
latent_dim = 512 |
|
|
|
model_path_s = hf_hub_download(repo_id="akhaliq/jojogan-stylegan2-ffhq-config-f", filename="stylegan2-ffhq-config-f.pt") |
|
original_generator = Generator(1024, latent_dim, 8, 2).to(device) |
|
ckpt = torch.load(model_path_s, map_location=lambda storage, loc: storage) |
|
original_generator.load_state_dict(ckpt["g_ema"], strict=False) |
|
mean_latent = original_generator.mean_latent(10000) |
|
|
|
|
|
|
|
generatorzombie = deepcopy(original_generator) |
|
generatorjojo = deepcopy(original_generator) |
|
|
|
transform = transforms.Compose( |
|
[ |
|
transforms.Resize((1024, 1024)), |
|
transforms.ToTensor(), |
|
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), |
|
] |
|
) |
|
|
|
|
|
|
|
modelzombie = hf_hub_download(repo_id="Awesimo/jojogan-zombie", filename="zombie.pt") |
|
ckptzombie = torch.load(modelzombie, map_location=lambda storage, loc: storage) |
|
|
|
generatorzombie.load_state_dict(ckptzombie, strict=False) |
|
|
|
|
|
modeljojo = hf_hub_download(repo_id="akhaliq/JoJoGAN-jojo", filename="jojo_preserve_color.pt") |
|
ckptjojo = torch.load(modeljojo, map_location=lambda storage, loc: storage) |
|
generatorjojo.load_state_dict(ckptjojo["g"], strict=False) |
|
|
|
def inference(img, model): |
|
img.save('out.jpg') |
|
aligned_face = align_face('out.jpg') |
|
|
|
my_w = projection(aligned_face, "test.pt", device).unsqueeze(0) |
|
if model == 'Zombie': |
|
with torch.no_grad(): |
|
my_sample = generatorzombie(my_w, input_is_latent=True) |
|
elif model == 'JoJo': |
|
with torch.no_grad(): |
|
my_sample = generatorjojo(my_w, input_is_latent=True) |
|
else: |
|
with torch.no_grad(): |
|
my_sample = generatorzombie(my_w, input_is_latent=True) |
|
|
|
|
|
npimage = my_sample[0].permute(1, 2, 0).detach().numpy() |
|
imageio.imwrite('filename.jpeg', npimage) |
|
return 'filename.jpeg' |
|
|
|
title = "JoJoGAN Test 🤖" |
|
examples=[['assets/samples/image01.jpg','Zombie'],['assets/samples/image02.jpg','JoJo'],['assets/samples/image03.jpg','Zombie'],['assets/samples/image04.jpg','JoJo']] |
|
gr.Interface(inference, [gr.inputs.Image(type="pil"),gr.inputs.Dropdown(choices=['Zombie', 'JoJo'], type="value", default='Zombie', label="Model")], gr.outputs.Image(type="file"),title=title,allow_flagging=False,examples=examples,allow_screenshot=False).launch() |
|
|
