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 time
 
 torch.set_grad_enabled(False)
 torch.set_num_threads(os.cpu_count())
@@ -240,6 +241,9 @@ class MinDalle:
             st.session_state.bar.progress(i/IMAGE_TOKEN_COUNT)
 
             torch.cuda.empty_cache()                
+            
+            time.sleep(0.5)
+            
             with torch.cuda.amp.autocast(dtype=self.dtype):
                 image_tokens[i + 1], attention_state = self.decoder.forward(
                     settings=settings,