Spaces:
Runtime error
Runtime error
import os | |
os.system("pip install gradio==2.8.0b5") | |
os.system("pip install -r requirements.txt") | |
os.system("pip freeze") | |
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 e4e_projection import projection as e4e_projection | |
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) | |
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) | |
generatorjojo = deepcopy(original_generator) | |
generatordisney = deepcopy(original_generator) | |
generatorjinx = deepcopy(original_generator) | |
generatorcaitlyn = deepcopy(original_generator) | |
generatoryasuho = deepcopy(original_generator) | |
generatorarcanemulti = deepcopy(original_generator) | |
generatorart = deepcopy(original_generator) | |
generatorspider = deepcopy(original_generator) | |
generatorsketch = deepcopy(original_generator) | |
generatordisneynoc = 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)), | |
] | |
) | |
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) | |
modeldisney = hf_hub_download(repo_id="akhaliq/jojogan-disney", filename="disney_preserve_color.pt") | |
ckptdisney = torch.load(modeldisney, map_location=lambda storage, loc: storage) | |
generatordisney.load_state_dict(ckptdisney["g"], strict=False) | |
modeljinx = hf_hub_download(repo_id="akhaliq/jojo-gan-jinx", filename="arcane_jinx_preserve_color.pt") | |
ckptjinx = torch.load(modeljinx, map_location=lambda storage, loc: storage) | |
generatorjinx.load_state_dict(ckptjinx["g"], strict=False) | |
modelcaitlyn = hf_hub_download(repo_id="akhaliq/jojogan-arcane", filename="arcane_caitlyn_preserve_color.pt") | |
ckptcaitlyn = torch.load(modelcaitlyn, map_location=lambda storage, loc: storage) | |
generatorcaitlyn.load_state_dict(ckptcaitlyn["g"], strict=False) | |
modelyasuho = hf_hub_download(repo_id="akhaliq/JoJoGAN-jojo", filename="jojo_yasuho_preserve_color.pt") | |
ckptyasuho = torch.load(modelyasuho, map_location=lambda storage, loc: storage) | |
generatoryasuho.load_state_dict(ckptyasuho["g"], strict=False) | |
model_arcane_multi = hf_hub_download(repo_id="akhaliq/jojogan-arcane", filename="arcane_multi_preserve_color.pt") | |
ckptarcanemulti = torch.load(model_arcane_multi, map_location=lambda storage, loc: storage) | |
generatorarcanemulti.load_state_dict(ckptarcanemulti["g"], strict=False) | |
modelart = hf_hub_download(repo_id="akhaliq/jojo-gan-art", filename="art.pt") | |
ckptart = torch.load(modelart, map_location=lambda storage, loc: storage) | |
generatorart.load_state_dict(ckptart["g"], strict=False) | |
modelSpiderverse = hf_hub_download(repo_id="akhaliq/jojo-gan-spiderverse", filename="Spiderverse-face-500iters-8face.pt") | |
ckptspider = torch.load(modelSpiderverse, map_location=lambda storage, loc: storage) | |
generatorspider.load_state_dict(ckptspider["g"], strict=False) | |
modelSketch = hf_hub_download(repo_id="akhaliq/jojogan-sketch", filename="sketch_multi.pt") | |
ckptsketch = torch.load(modelSketch, map_location=lambda storage, loc: storage) | |
generatorsketch.load_state_dict(ckptsketch["g"], strict=False) | |
modeldisneynoc = hf_hub_download(repo_id="algomuffin/disney", filename="disney.pt") | |
ckptMy = torch.load(modeldisneynoc, map_location=lambda storage, loc: storage) | |
generatordisneynoc.load_state_dict(ckptMy["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 == 'JoJo': | |
with torch.no_grad(): | |
my_sample = generatorjojo(my_w, input_is_latent=True) | |
elif model == 'Disney': | |
with torch.no_grad(): | |
my_sample = generatordisneynoc(my_w, input_is_latent=True) | |
elif model == 'Jinx': | |
with torch.no_grad(): | |
my_sample = generatorjinx(my_w, input_is_latent=True) | |
elif model == 'Caitlyn': | |
with torch.no_grad(): | |
my_sample = generatorcaitlyn(my_w, input_is_latent=True) | |
elif model == 'Yasuho': | |
with torch.no_grad(): | |
my_sample = generatoryasuho(my_w, input_is_latent=True) | |
elif model == 'Arcane Multi': | |
with torch.no_grad(): | |
my_sample = generatorarcanemulti(my_w, input_is_latent=True) | |
elif model == 'Art': | |
with torch.no_grad(): | |
my_sample = generatorart(my_w, input_is_latent=True) | |
elif model == 'Spider-Verse': | |
with torch.no_grad(): | |
my_sample = generatorspider(my_w, input_is_latent=True) | |
else: | |
with torch.no_grad(): | |
my_sample = generatorsketch(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_fork" | |
description = "Gradio Demo for JoJoGAN: This is a fork made by algomuffin in order to test the network with specific parameters. To use it, simply upload your image, or click one of the examples to load them." | |
article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2112.11641' target='_blank'>JoJoGAN: One Shot Face Stylization</a>| <a href='https://github.com/mchong6/JoJoGAN' target='_blank'>Github Repo Pytorch</a></p> <center><img src='https://visitor-badge.glitch.me/badge?page_id=akhaliq_jojogan' alt='visitor badge'></center>" | |
examples=[['mona.png','Disney']] | |
gr.Interface(inference, [gr.inputs.Image(type="pil"),gr.inputs.Dropdown(choices=['Disney'], type="value", default='Disney', label="Model")], gr.outputs.Image(type="file"),title=title,description=description,article=article,allow_flagging=False,examples=examples,allow_screenshot=False).launch(enable_queue=True, cache_examples=True) | |