Remove unused models

app.py

@@ -69,15 +69,6 @@ mean_latent = original_generator.mean_latent(10000)
 #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)
-
 
 transform = transforms.Compose(
     [
@@ -89,7 +80,7 @@ transform = transforms.Compose(
 
 
 #ZOMBIE
-modelzombie = hf_hub_download(repo_id="Awesimo/jojogan-zombie", filename="jojo_preserve_color.pt")
+modelzombie = hf_hub_download(repo_id="Awesimo/jojogan-zombie", filename="jojo_zombie_preserve_color.pt")
 ckptzombie = torch.load(modelzombie, map_location=lambda storage, loc: storage)
 generatorzombie.load_state_dict(ckptzombie["g"], strict=False)
 
@@ -98,47 +89,6 @@ modeljojo = hf_hub_download(repo_id="akhaliq/JoJoGAN-jojo", filename="jojo_prese
 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')
@@ -149,31 +99,10 @@ def inference(img, model):
             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)
-    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)
+            my_sample = generatorjojo(my_w, input_is_latent=True)
     else:
         with torch.no_grad():
-            my_sample = generatorsketch(my_w, input_is_latent=True)
+            my_sample = generatorzombie(my_w, input_is_latent=True)
 
 
     npimage = my_sample[0].permute(1, 2, 0).detach().numpy()
@@ -181,5 +110,5 @@ def inference(img, model):
     return 'filename.jpeg'
 
 title = "JoJoGAN Test 🤖"
-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=['Zombie', '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()
+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()