Fix

app.py

@@ -2,10 +2,7 @@ import os
 from PIL import Image
 import torch
 import gradio as gr
-import torch
 torch.backends.cudnn.benchmark = True
-from torchvision import transforms, utils
-from PIL import Image
 import math
 import random
 import numpy as np
@@ -18,11 +15,8 @@ import time
 from copy import deepcopy
 import imageio
 
-import os
 import sys
-import numpy as np
 from PIL import Image
-import torch
 import torchvision.transforms as transforms
 from argparse import Namespace
 from e4e.utils.common import tensor2im
@@ -114,9 +108,9 @@ def run_alignment(image_path):
 def gen_im(ffhq_codes, dog_codes, cat_codes, model_type='ffhq'):
   if model_type=='ffhq':
     imgs, _ = ffhq_decoder([ffhq_codes], input_is_latent=True, randomize_noise=False, return_latents=True)
-  elif model_type=='dog':
+  elif model_type=='Dog':
     imgs, _ = dog_decoder([dog_codes], input_is_latent=True, randomize_noise=False, return_latents=True)
-  elif model_type=='cat':
+  elif model_type=='Cat':
     imgs, _ = cat_decoder([cat_codes], input_is_latent=True, randomize_noise=False, return_latents=True)
   else:
     imgs, _ = custom_decoder([custom_codes], input_is_latent=True, randomize_noise=False, return_latents=True)
@@ -125,7 +119,7 @@ def gen_im(ffhq_codes, dog_codes, cat_codes, model_type='ffhq'):
 def set_seed(rd):
     torch.manual_seed(rd)
 
-def inference(img):  
+def inference(img, model):  
     random_seed = round(time.time() * 1000)
     set_seed(random_seed)
     
@@ -143,17 +137,17 @@ def inference(img):
     dog_codes = dog_encoder(transformed_image.unsqueeze(0).to(device).float())
     dog_codes = dog_codes + ffhq_latent_avg.repeat(dog_codes.shape[0], 1, 1)
         
-    animal = "cat"
-    npimage = gen_im(ffhq_codes, dog_codes, cat_codes, animal)
+    npimage = gen_im(ffhq_codes, dog_codes, cat_codes, model)
 	
     imageio.imwrite('filename.jpeg', npimage)
     return 'filename.jpeg'
   
 title = "PetBreeder v1.1"
-description = "Gradio Demo for PetBreeder."
+description = "Gradio Demo for PetBreeder. Based on [Colab](https://colab.research.google.com/github/tg-bomze/collection-of-notebooks/blob/master/PetBreeder.ipynb) by [@MLArt](https://t.me/MLArt)."
 
 gr.Interface(inference, 
-[gr.inputs.Image(type="pil")], 
+[gr.inputs.Image(type="pil"),
+gr.inputs.Dropdown(choices=['Cat','Dog'], type='value', default='Cat', label='Model')]
 gr.outputs.Image(type="file"),
 title=title,
 description=description).launch()