Add zombie model

app.py

@@ -15,10 +15,6 @@ 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
 
@@ -44,7 +40,6 @@ net = pSp(opts, device).eval().to(device)
 @ torch.no_grad()
 def projection(img, name, device='cuda'):
 
-
     transform = transforms.Compose(
         [
             transforms.Resize(256),
@@ -60,12 +55,8 @@ def projection(img, name, device='cuda'):
     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")
@@ -74,22 +65,17 @@ 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)
 
+#MODELS
+generatorzombie = deepcopy(original_generator)
+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)
 
 
@@ -102,69 +88,68 @@ transform = transforms.Compose(
 )
 
 
+#ZOMBIE
+modelzombie = hf_hub_download(repo_id="akhaliq/JoJoGAN-jojo", filename="jojo_preserve_color.pt")
+ckptzombie = torch.load(modelzombie, map_location=lambda storage, loc: storage)
+generatorzombie.load_state_dict(ckptzombie["g"], strict=False)
 
-
+#JOJO
 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)
 
-
+#DISNEY
 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)
 
-
+#JINX
 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)
 
-
+#CAITLYN
 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)
 
-
+#YASHUO
 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)
 
-
+#ARCANE
 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)
 
-
+#ART
 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)
 
-
+#SPIDER
 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)
 
+#SKETCH
 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':
+    if model == 'Zombie':
         with torch.no_grad():
-            my_sample = generatorjojo(my_w, input_is_latent=True)
+            my_sample = generatorzombie(my_w, input_is_latent=True)
+    elif model == 'JoJo':
+        with torch.no_grad():
+            my_sample = generatordisney(my_w, input_is_latent=True)
     elif model == 'Disney':
         with torch.no_grad():
             my_sample = generatordisney(my_w, input_is_latent=True)
@@ -196,5 +181,5 @@ def inference(img, model):
     return 'filename.jpeg'
 
 title = "JoJoGAN Test 🤖"
-examples=[['assets/image01.jpg','JoJo'],['assets/image02.jpg','Disney'],['assets/image03.jpg','Jinx'],['assets/image04.jpg','Sketch']]
-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,allow_flagging=False,examples=examples,allow_screenshot=False).launch()
+examples=[['assets/samples/image01.jpg','JoJo'],['assets/samples/image02.jpg','Disney'],['assets/samples/image03.jpg','Jinx'],['assets/samples/image04.jpg','Sketch']]
+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='Zombie', label="Model")], gr.outputs.Image(type="file"),title=title,allow_flagging=False,examples=examples,allow_screenshot=False).launch()