Update app.py

app.py

@@ -34,7 +34,7 @@ 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")
+model_path_e = hf_hub_download(repo_id="Abhinowww/Capstone", filename="e4e_ffhq_encode.pt")
 ckpt = torch.load(model_path_e, map_location='cpu')
 opts = ckpt['opts']
 opts['checkpoint_path'] = model_path_e
@@ -68,29 +68,29 @@ 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")
+model_path_s = hf_hub_download(repo_id="Abhinowww/Capstone", 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)
+generatorjoker = deepcopy(original_generator)
 
-generatordisney = deepcopy(original_generator)
+generatorvoldemort = deepcopy(original_generator)
 
-generatorjinx = deepcopy(original_generator)
+# generatorjinx = deepcopy(original_generator)
 
-generatorcaitlyn = deepcopy(original_generator)
+# generatorcaitlyn = deepcopy(original_generator)
 
-generatoryasuho = deepcopy(original_generator)
+# generatoryasuho = deepcopy(original_generator)
 
-generatorarcanemulti = deepcopy(original_generator)
+# generatorarcanemulti = deepcopy(original_generator)
 
-generatorart = deepcopy(original_generator)
+# generatorart = deepcopy(original_generator)
 
-generatorspider = deepcopy(original_generator)
+# generatorspider = deepcopy(original_generator)
 
-generatorsketch = deepcopy(original_generator)
+# generatorsketch = deepcopy(original_generator)
 
 
 transform = transforms.Compose(
@@ -104,58 +104,58 @@ transform = transforms.Compose(
 
 
 
-modeljojo = hf_hub_download(repo_id="akhaliq/JoJoGAN-jojo", filename="jojo_preserve_color.pt")
+modeljoker = hf_hub_download(repo_id="Abhinowww/Capstone", filename="JokerEightHundredFalse.pt")
 
 
-ckptjojo = torch.load(modeljojo, map_location=lambda storage, loc: storage)
-generatorjojo.load_state_dict(ckptjojo["g"], strict=False)
+ckptjoker = torch.load(modeljoker, map_location=lambda storage, loc: storage)
+generatorjoker.load_state_dict(ckptjojo["g"], strict=False)
 
 
-modeldisney = hf_hub_download(repo_id="akhaliq/jojogan-disney", filename="disney_preserve_color.pt")
+modelvoldemort = hf_hub_download(repo_id="Abhinowww/Capstone", filename="VoldemortEightHundredFalse.pt")
 
-ckptdisney = torch.load(modeldisney, map_location=lambda storage, loc: storage)
+ckptdisney = torch.load(modelvoldemort, 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")
+# 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)
+# 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")
+# 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)
+# 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")
+# 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)
+# 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")
+# 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)
+# 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")
+# 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)
+# 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")
+# 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)
+# 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")
+# 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)
+# 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')  
@@ -168,37 +168,38 @@ def inference(img, model):
     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)
+    # 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."
+title = "Capstone"
+description = "Capstone Project. 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="pil"),title=title,description=description,article=article,allow_flagging=False,examples=examples,allow_screenshot=False).launch()
+examples=[['mona.png','Joker']]
+# 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="pil"),title=title,description=description,article=article,allow_flagging=False,examples=examples,allow_screenshot=False).launch()
+gr.Interface(inference, [gr.inputs.Image(type="pil"),gr.inputs.Dropdown(choices=['Joker', 'Voldemort'], type="value", default='Joker', label="Model")], gr.outputs.Image(type="pil"),title=title,description=description,article=article,allow_flagging=False,examples=examples,allow_screenshot=False).launch()