Update app.py

app.py

@@ -123,24 +123,30 @@ def inference(img, model):
     random_seed = round(time.time() * 1000)
     set_seed(random_seed)
     
-    img.save('out.jpg')  
-
-    input_image = run_alignment('out.jpg')
-    transformed_image = transform(input_image)
-	
-    ffhq_codes = ffhq_encoder(transformed_image.unsqueeze(0).to(device).float())
-    ffhq_codes = ffhq_codes + ffhq_latent_avg.repeat(ffhq_codes.shape[0], 1, 1)
-
-    cat_codes = cat_encoder(transformed_image.unsqueeze(0).to(device).float())
-    cat_codes = cat_codes + cat_latent_avg.repeat(cat_codes.shape[0], 1, 1)
-
-    dog_codes = dog_encoder(transformed_image.unsqueeze(0).to(device).float())
-    dog_codes = dog_codes + dog_latent_avg.repeat(dog_codes.shape[0], 1, 1)
+    try:
+        img.save('out.jpg')  
         
-    npimage = gen_im(ffhq_codes, dog_codes, cat_codes, model)
-	
-    imageio.imwrite('filename.jpeg', npimage)
-    return 'filename.jpeg'
+        try:
+            input_image = run_alignment('out.jpg')
+        except:
+            return 'out.jpg'
+        transformed_image = transform(input_image)
+    	
+        ffhq_codes = ffhq_encoder(transformed_image.unsqueeze(0).to(device).float())
+        ffhq_codes = ffhq_codes + ffhq_latent_avg.repeat(ffhq_codes.shape[0], 1, 1)
+    
+        cat_codes = cat_encoder(transformed_image.unsqueeze(0).to(device).float())
+        cat_codes = cat_codes + cat_latent_avg.repeat(cat_codes.shape[0], 1, 1)
+    
+        dog_codes = dog_encoder(transformed_image.unsqueeze(0).to(device).float())
+        dog_codes = dog_codes + dog_latent_avg.repeat(dog_codes.shape[0], 1, 1)
+            
+        npimage = gen_im(ffhq_codes, dog_codes, cat_codes, model)
+    	
+        imageio.imwrite('filename.jpeg', npimage)
+        return 'filename.jpeg'
+    except:
+        pass
   
 title = "PetBreeder v1.1"
 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)."