Update min_dalle/min_dalle.py

min_dalle/min_dalle.py

@@ -12,6 +12,8 @@ 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())
 torch.backends.cudnn.enabled = True
@@ -235,7 +237,11 @@ class MinDalle:
             dtype=torch.float32,
             device=self.device
         )
+        
+        tracemalloc.start()
+        
         for i in range( IMAGE_TOKEN_COUNT ):
+        
             if(st.session_state.page != 0):
                 break
             st.session_state.bar.progress(i/IMAGE_TOKEN_COUNT)
@@ -262,7 +268,12 @@ class MinDalle:
                         is_seamless=is_seamless,
                         is_verbose=is_verbose
                     )
-                    
+            # displaying the memory
+            print(tracemalloc.get_traced_memory())
+ 
+        # stopping the library
+        tracemalloc.stop()
+