Update min_dalle/min_dalle.py

min_dalle/min_dalle.py

@@ -10,6 +10,7 @@ from typing import Iterator
 from .text_tokenizer import TextTokenizer
 from .models import DalleBartEncoder, DalleBartDecoder, VQGanDetokenizer
 import streamlit as st
+import gc
 
 torch.set_grad_enabled(False)
 torch.set_num_threads(os.cpu_count())
@@ -235,12 +236,14 @@ class MinDalle:
             device=self.device
         )
         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() 
             #torch.cpu.empty_cache()                
+            gc.collect()
         
             image_tokens[i + 1], attention_state = self.decoder.forward(
                 settings=settings,
@@ -257,7 +260,7 @@ class MinDalle:
                     is_seamless=is_seamless,
                     is_verbose=is_verbose
                 )
-            del attention_state
+            
 
     def generate_image_stream(self, *args, **kwargs) -> Iterator[Image.Image]:
         image_stream = self.generate_raw_image_stream(*args, **kwargs)