Update pipeline.py

pipeline.py

@@ -2,6 +2,7 @@ import torch
 import nltk
 import io
 import base64
+import shutil
 from torchvision import transforms
 
 from pytorch_pretrained_biggan import BigGAN, one_hot_from_names, truncated_noise_sample
@@ -11,8 +12,13 @@ class PreTrainedPipeline():
         """
         Initialize model
         """
-        nltk.download('wordnet')
-        self.model = BigGAN.from_pretrained(path)
+        try:
+            self.model = BigGAN.from_pretrained(path)
+        except (IOError, OSError):
+            directory = "/data/corpora"
+            shutil.rmtree(directory)
+            nltk.download('wordnet')
+            self.model = BigGAN.from_pretrained(path)
         self.truncation = 0.1
 
     def __call__(self, inputs: str):