Update min_dalle/min_dalle.py

min_dalle/min_dalle.py

@@ -242,8 +242,6 @@ class MinDalle:
 
             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,
@@ -253,8 +251,7 @@ class MinDalle:
                     prev_tokens=image_tokens[i],
                     token_index=token_indices[[i]]
                 )
-                
-            #torch.cuda.empty_cache()
+             
 
             with torch.cuda.amp.autocast(dtype=torch.float16):
                 if ((i + 1) % 16 == 0 and progressive_outputs) or i + 1 == 256:
@@ -263,6 +260,10 @@ class MinDalle:
                         is_seamless=is_seamless,
                         is_verbose=is_verbose
                     )
+            del attention_state
+            del image_tokens
+                    
+            
 
     def generate_image_stream(self, *args, **kwargs) -> Iterator[Image.Image]:
         image_stream = self.generate_raw_image_stream(*args, **kwargs)