Update min_dalle/min_dalle.py

min_dalle/min_dalle.py

@@ -17,7 +17,7 @@ torch.backends.cudnn.enabled = True
 torch.backends.cudnn.allow_tf32 = True
 
 MIN_DALLE_REPO = 'https://huggingface.co/kuprel/min-dalle/resolve/main/'
-IMAGE_TOKEN_COUNT = 240
+IMAGE_TOKEN_COUNT = 256
 
 
 class MinDalle:
@@ -239,28 +239,27 @@ class MinDalle:
                 break
             st.session_state.bar.progress(i/IMAGE_TOKEN_COUNT)
             
-            try:
-                #torch.cuda.empty_cache() 
-                #torch.cpu.empty_cache()                
-                #with torch.cuda.amp.autocast(dtype=self.dtype):
-                image_tokens[i + 1], attention_state = self.decoder.forward(
-                    settings=settings,
-                    attention_mask=attention_mask,
-                    encoder_state=encoder_state,
-                    attention_state=attention_state,
-                    prev_tokens=image_tokens[i],
-                    token_index=token_indices[[i]]
+      
+            #torch.cuda.empty_cache() 
+            #torch.cpu.empty_cache()                
+            #with torch.cuda.amp.autocast(dtype=self.dtype):
+            image_tokens[i + 1], attention_state = self.decoder.forward(
+                settings=settings,
+                attention_mask=attention_mask,
+                encoder_state=encoder_state,
+                attention_state=attention_state,
+                prev_tokens=image_tokens[i],
+                token_index=token_indices[[i]]
+            )
+
+            # 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
                 )
-    
-               # 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
-                    )
-            except Exception as e:
-                print(e)
+           
 
     def generate_image_stream(self, *args, **kwargs) -> Iterator[Image.Image]:
         image_stream = self.generate_raw_image_stream(*args, **kwargs)
@@ -277,6 +276,7 @@ class MinDalle:
             image = image.transpose(1, 0)
             image = image.reshape([grid_size ** 2, 2 ** 8, 2 ** 8, 3])
             yield image
+            del image
 
 
     def generate_image(self, *args, **kwargs) -> Image.Image: