Update min_dalle/min_dalle.py

min_dalle/min_dalle.py

@@ -66,7 +66,7 @@ class MinDalle:
             self.init_decoder()
             self.init_detokenizer()
 
-
+    @st.cache
     def download_tokenizer(self):
         if self.is_verbose: print("downloading tokenizer params")
         suffix = '' if self.is_mega else '_mini'
@@ -76,27 +76,27 @@ class MinDalle:
         with open(self.vocab_path, 'wb') as f: f.write(vocab.content)
         with open(self.merges_path, 'wb') as f: f.write(merges.content)
 
-
+    @st.cache
     def download_encoder(self):
         if self.is_verbose: print("downloading encoder params")
         suffix = '' if self.is_mega else '_mini'
         params = requests.get(MIN_DALLE_REPO + 'encoder{}.pt'.format(suffix))
         with open(self.encoder_params_path, 'wb') as f: f.write(params.content)
 
-
+    @st.cache
     def download_decoder(self):
         if self.is_verbose: print("downloading decoder params")
         suffix = '' if self.is_mega else '_mini'
         params = requests.get(MIN_DALLE_REPO + 'decoder{}.pt'.format(suffix))
         with open(self.decoder_params_path, 'wb') as f: f.write(params.content)
     
-
+    @st.cache
     def download_detokenizer(self):
         if self.is_verbose: print("downloading detokenizer params")
         params = requests.get(MIN_DALLE_REPO + 'detoker.pt')
         with open(self.detoker_params_path, 'wb') as f: f.write(params.content)
 
-
+    @st.cache
     def init_tokenizer(self):
         is_downloaded = os.path.exists(self.vocab_path)
         is_downloaded &= os.path.exists(self.merges_path)
@@ -108,7 +108,7 @@ class MinDalle:
             merges = f.read().split("\n")[1:-1]
         self.tokenizer = TextTokenizer(vocab, merges)
 
-
+    @st.cache
     def init_encoder(self):
         is_downloaded = os.path.exists(self.encoder_params_path)
         if not is_downloaded: self.download_encoder()
@@ -127,7 +127,7 @@ class MinDalle:
         del params
         self.encoder = self.encoder.to(device=self.device)
 
-
+    @st.cache
     def init_decoder(self):
         is_downloaded = os.path.exists(self.decoder_params_path)
         if not is_downloaded: self.download_decoder()
@@ -145,7 +145,7 @@ class MinDalle:
         del params
         self.decoder = self.decoder.to(device=self.device)
 
-
+    @st.cache
     def init_detokenizer(self):
         is_downloaded = os.path.exists(self.detoker_params_path)
         if not is_downloaded: self.download_detokenizer()