Update min_dalle/min_dalle.py

min_dalle/min_dalle.py

@@ -24,6 +24,7 @@ IMAGE_TOKEN_COUNT = 256
 
 
 class MinDalle:
+    @st.cache
     def __init__(
         self,
         models_root: str = 'pretrained',
@@ -66,7 +67,6 @@ 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,23 @@ 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 +104,6 @@ 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()