add caitlyn model

app.py

@@ -4,12 +4,6 @@ import torch
 import gradio as gr
 os.system("pip install gradio==2.5.3")
 
-#os.system("pip install facexlib")
-
-#from facexlib.utils.face_restoration_helper import FaceRestoreHelper
-#os.system("pip install autocrop")
-#os.system("pip install dlib")
-#from autocrop import Cropper
 import torch
 torch.backends.cudnn.benchmark = True
 from torchvision import transforms, utils
@@ -38,15 +32,6 @@ os.system("wget http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2"
 os.system("bzip2 -dk shape_predictor_68_face_landmarks.dat.bz2")
 os.system("mv shape_predictor_68_face_landmarks.dat models/dlibshape_predictor_68_face_landmarks.dat")
 
-#cropper = Cropper(face_percent=80)
-
-#face_helper = FaceRestoreHelper(
-        #upscale_factor=0,
-        #face_size=512,
-        #crop_ratio=(1, 1),
-        #det_model='retinaface_resnet50',
-        #save_ext='png',
-        #device='cpu')
 
 device = 'cpu' 
 
@@ -54,19 +39,19 @@ os.system("gdown https://drive.google.com/uc?id=1_cTsjqzD_X9DK3t3IZE53huKgnzj_bt
 
 latent_dim = 512
 
-# Load original generator
 original_generator = Generator(1024, latent_dim, 8, 2).to(device)
 ckpt = torch.load('stylegan2-ffhq-config-f.pt', map_location=lambda storage, loc: storage)
 original_generator.load_state_dict(ckpt["g_ema"], strict=False)
 mean_latent = original_generator.mean_latent(10000)
 
-# to be finetuned generator
 generatorjojo = deepcopy(original_generator)
 
 generatordisney = deepcopy(original_generator)
 
 generatorjinx = deepcopy(original_generator)
 
+generatorcaitlyn = deepcopy(original_generator)
+
 
 
 transform = transforms.Compose(
@@ -95,22 +80,15 @@ os.system("gdown https://drive.google.com/uc?id=1jElwHxaYPod5Itdy18izJk49K1nl4ne
 ckptjinx = torch.load('arcane_jinx_preserve_color.pt', map_location=lambda storage, loc: storage)
 generatorjinx.load_state_dict(ckptjinx["g"], strict=False)
 
+os.system("gdown https://drive.google.com/uc?id=1cUTyjU-q98P75a8THCaO545RTwpVV-aH")
+
+ckptcaitlyn = torch.load('arcane_caitlyn_preserve_color.pt', map_location=lambda storage, loc: storage)
+generatorcaitlyn.load_state_dict(ckptcaitlyn["g"], strict=False)
+
 
 def inference(img, model):    
-    #face_helper.clean_all()
     aligned_face = align_face(img)
-    #cropped_array = cropper.crop(img[:,:,::-1])
-    
-    #if cropped_array.any():
-        #aligned_face = Image.fromarray(cropped_array)
-    #else:
-        #aligned_face = Image.fromarray(img[:,:,::-1])
-    
-    #face_helper.read_image(img)
-    #face_helper.get_face_landmarks_5(only_center_face=False, eye_dist_threshold=10)
-    #face_helper.align_warp_face(save_cropped_path="/home/user/app/")
-    #pilimg = Image.open("/home/user/app/_02.png")
-    
+        
     my_w = e4e_projection(aligned_face, "test.pt", device).unsqueeze(0)
     if model == 'JoJo':
         with torch.no_grad():
@@ -118,9 +96,12 @@ def inference(img, model):
     elif model == 'Disney':
         with torch.no_grad():
             my_sample = generatordisney(my_w, input_is_latent=True)
-    else:
+    elif model == 'Jinx':
         with torch.no_grad():
             my_sample = generatorjinx(my_w, input_is_latent=True)
+    else:
+        with torch.no_grad():
+            my_sample = generatorcaitlyn(my_w, input_is_latent=True)
             
     
     npimage = my_sample[0].permute(1, 2, 0).detach().numpy()
@@ -133,4 +114,4 @@ 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','Jinx']]
-gr.Interface(inference, [gr.inputs.Image(type="filepath"),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()
+gr.Interface(inference, [gr.inputs.Image(type="filepath"),gr.inputs.Dropdown(choices=['JoJo', 'Disney','Jinx','Caitlyn'], 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()