Create app.py

app.py

@@ -0,0 +1,204 @@
+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 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)
+
+@ 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)
+
+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)
+
+
+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)
+
+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 = generatordisney(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"
+description = "Gradio Demo for JoJoGAN: One Shot Face Stylization. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
+
+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','Jinx']]
+gr.Interface(inference, [gr.inputs.Image(type="pil"),gr.inputs.Dropdown(choices=['JoJo', 'Disney','Jinx','Caitlyn','Yasuho','Arcane Multi','Art','Spider-Verse','Sketch'], type="value", default='JoJo', label="Model")], gr.outputs.Image(type="file"),title=title,description=description,article=article,allow_flagging=False,examples=examples,allow_screenshot=False).launch()