Update min_dalle/min_dalle.py

min_dalle/min_dalle.py

@@ -10,9 +10,6 @@ from typing import Iterator
 from .text_tokenizer import TextTokenizer
 from .models import DalleBartEncoder, DalleBartDecoder, VQGanDetokenizer
 import streamlit as st
-import time
-
-import tracemalloc
 
 torch.set_grad_enabled(False)
 torch.set_num_threads(os.cpu_count())
@@ -24,7 +21,6 @@ IMAGE_TOKEN_COUNT = 256
 
 
 class MinDalle:
-    @st.cache
     def __init__(
         self,
         models_root: str = 'pretrained',
@@ -67,6 +63,7 @@ class MinDalle:
             self.init_decoder()
             self.init_detokenizer()
 
+
     def download_tokenizer(self):
         if self.is_verbose: print("downloading tokenizer params")
         suffix = '' if self.is_mega else '_mini'
@@ -76,23 +73,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)
 
+
     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)
 
+
     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)
     
+
     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)
 
+
     def init_tokenizer(self):
         is_downloaded = os.path.exists(self.vocab_path)
         is_downloaded &= os.path.exists(self.merges_path)
@@ -104,6 +105,7 @@ class MinDalle:
             merges = f.read().split("\n")[1:-1]
         self.tokenizer = TextTokenizer(vocab, merges)
 
+
     def init_encoder(self):
         is_downloaded = os.path.exists(self.encoder_params_path)
         if not is_downloaded: self.download_encoder()
@@ -122,6 +124,7 @@ class MinDalle:
         del params
         self.encoder = self.encoder.to(device=self.device)
 
+
     def init_decoder(self):
         is_downloaded = os.path.exists(self.decoder_params_path)
         if not is_downloaded: self.download_decoder()
@@ -138,7 +141,8 @@ class MinDalle:
         self.decoder.load_state_dict(params, strict=False)
         del params
         self.decoder = self.decoder.to(device=self.device)
-        
+
+
     def init_detokenizer(self):
         is_downloaded = os.path.exists(self.detoker_params_path)
         if not is_downloaded: self.download_detokenizer()
@@ -230,17 +234,12 @@ class MinDalle:
             dtype=torch.float32,
             device=self.device
         )
-        
-        tracemalloc.start()
-        
-        for i in range( IMAGE_TOKEN_COUNT ):
-        
+        for i in range(IMAGE_TOKEN_COUNT):
             if(st.session_state.page != 0):
                 break
             st.session_state.bar.progress(i/IMAGE_TOKEN_COUNT)
 
             torch.cuda.empty_cache()                
-            
             with torch.cuda.amp.autocast(dtype=self.dtype):
                 image_tokens[i + 1], attention_state = self.decoder.forward(
                     settings=settings,
@@ -250,27 +249,14 @@ class MinDalle:
                     prev_tokens=image_tokens[i],
                     token_index=token_indices[[i]]
                 )
-                
-      
-             
 
-            with torch.cuda.amp.autocast(dtype=torch.float16):
-                if ((i + 1) % 16 == 0 and progressive_outputs) or i + 1 == 256:
+            with torch.cuda.amp.autocast(dtype=torch.float32):
+                if ((i + 1) % 32 == 0 and progressive_outputs) or i + 1 == 256:
                     yield self.image_grid_from_tokens(
                         image_tokens=image_tokens[1:].T,
                         is_seamless=is_seamless,
                         is_verbose=is_verbose
                     )
-                    
-            # displaying the memory
-            print(tracemalloc.get_traced_memory())
- 
-        # stopping the library
-        tracemalloc.stop()
-                            
-
-                    
-            
 
     def generate_image_stream(self, *args, **kwargs) -> Iterator[Image.Image]:
         image_stream = self.generate_raw_image_stream(*args, **kwargs)