Update app.py

app.py

@@ -102,6 +102,8 @@ generatorart = deepcopy(original_generator)
 
 generatorspider = deepcopy(original_generator)
 
+generatorsketch = deepcopy(original_generator)
+
 
 transform = transforms.Compose(
     [
@@ -162,6 +164,10 @@ modelSpiderverse = hf_hub_download(repo_id="akhaliq/jojo-gan-spiderverse", filen
 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/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')  
@@ -189,9 +195,12 @@ def inference(img, model):
     elif model == 'Art':
         with torch.no_grad():
             my_sample = generatorart(my_w, input_is_latent=True)
-    else:
+    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()
@@ -204,4 +213,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>"
 
 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'], type="value", default='JoJo', label="Model")], gr.outputs.Image(type="file"),title=title,description=description,article=article,allow_flagging=False,examples=examples,allow_screenshot=False).launch(enable_queue=True, cache_examples=True)
+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,description=description,article=article,allow_flagging=False,examples=examples,allow_screenshot=False).launch(enable_queue=True, cache_examples=True)