Update app.py

app.py

@@ -50,6 +50,9 @@ generatorjojo = deepcopy(original_generator)
 
 generatordisney = deepcopy(original_generator)
 
+generatorjinx = deepcopy(original_generator)
+
+
 
 transform = transforms.Compose(
     [
@@ -72,19 +75,25 @@ os.system("gdown https://drive.google.com/uc?id=1Bnh02DjfvN_Wm8c4JdOiNV4q9J7Z_ts
 ckptdisney = torch.load('disney_preserve_color.pt', map_location=lambda storage, loc: storage)
 generatordisney.load_state_dict(ckptdisney["g"], strict=False)
 
+os.system("gdown https://drive.google.com/uc?id=1jElwHxaYPod5Itdy18izJk49K1nl4ney")
+
+ckptjinx = torch.load('arcane_jinx_preserve_color.pt', map_location=lambda storage, loc: storage)
+generatorjinx.load_state_dict(ckptjinx["g"], strict=False)
+
 
 def inference(img, model):    
     aligned_face = align_face(img)
-
     my_w = e4e_projection(aligned_face, "test.pt", device).unsqueeze(0)
     if model == 'JoJo':
         with torch.no_grad():
-        
             my_sample = generatorjojo(my_w, input_is_latent=True)  
-    else:
+    elif model == 'Disney':
         with torch.no_grad():
-        
             my_sample = generatordisney(my_w, input_is_latent=True)
+    else:
+        with torch.no_grad():
+            my_sample = generatorjinx(my_w, input_is_latent=True)
+            
     
     npimage = my_sample[0].permute(1, 2, 0).detach().numpy()
     imageio.imwrite('filename.jpeg', npimage)
@@ -95,5 +104,5 @@ description = "Gradio Demo for JoJoGAN: One Shot Face Stylization. To use it, si
 
 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> <p style='text-align: center'>samples from repo: <img src='https://raw.githubusercontent.com/mchong6/JoJoGAN/main/teaser.jpg' alt='animation'/></p>"
 
-examples=[['iu.jpeg']]
-gr.Interface(inference, [gr.inputs.Image(type="filepath",shape=(256,256)),gr.inputs.Dropdown(choices=['JoJo', 'Disney'], type="value", default='JoJo', label="Model")], gr.outputs.Image(type="file"),title=title,description=description,article=article,enable_queue=True,allow_flagging=False,examples=examples).launch()
+examples=[['iu.jpeg','Jinx']]
+gr.Interface(inference, [gr.inputs.Image(type="filepath",shape=(256,256)),gr.inputs.Dropdown(choices=['JoJo', 'Disney','Jinx'], type="value", default='JoJo', label="Model")], gr.outputs.Image(type="file"),title=title,description=description,article=article,enable_queue=True,allow_flagging=False,examples=examples).launch()