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 = 128
+IMAGE_TOKEN_COUNT = 256
 
 
 class MinDalle:
@@ -177,7 +177,7 @@ class MinDalle:
         progressive_outputs: bool = False,
         is_seamless: bool = False,
         temperature: float = 1,
-        top_k: int = 128,
+        top_k: int = 256,
         supercondition_factor: int = 16,
         is_verbose: bool = False
     ) -> Iterator[FloatTensor]:
@@ -239,8 +239,8 @@ class MinDalle:
                 break
             st.session_state.bar.progress(i/IMAGE_TOKEN_COUNT)
 
-            #torch.cuda.empty_cache()   
-            #torch.device('cpu').empty_cache()            
+            #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,
@@ -252,7 +252,7 @@ class MinDalle:
                 )
 
             with torch.cuda.amp.autocast(dtype=torch.float32):
-                if ((i + 1) % 32 == 0 and progressive_outputs) or i + 1 == 128:
+                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,
@@ -270,7 +270,7 @@ class MinDalle:
         image_stream = self.generate_raw_image_stream(*args, **kwargs)
         for image in image_stream:
             grid_size = kwargs['grid_size']
-            image = image.view([grid_size * 128, grid_size, 128, 3])
+            image = image.view([grid_size * 256, grid_size, 256, 3])
             image = image.transpose(1, 0)
             image = image.reshape([grid_size ** 2, 2 ** 8, 2 ** 8, 3])
             yield image