Update min_dalle/min_dalle.py

min_dalle/min_dalle.py

@@ -177,7 +177,7 @@ class MinDalle:
         progressive_outputs: bool = False,
         is_seamless: bool = False,
         temperature: float = 1,
-        top_k: int = 256,
+        top_k: int = 128,
         supercondition_factor: int = 16,
         is_verbose: bool = False
     ) -> Iterator[FloatTensor]:
@@ -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 == 256:
+                if ((i + 1) % 32 == 0 and progressive_outputs) or i + 1 == 128:
                     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 * 256, grid_size, 256, 3])
+            image = image.view([grid_size * 128, grid_size, 128, 3])
             image = image.transpose(1, 0)
             image = image.reshape([grid_size ** 2, 2 ** 8, 2 ** 8, 3])
             yield image