added grid_size parameter to generate a grid of images

image_from_text.py

@@ -11,9 +11,10 @@ parser.add_argument('--no-mega', dest='mega', action='store_false')
 parser.set_defaults(mega=False)
 parser.add_argument('--text', type=str, default='alien life')
 parser.add_argument('--seed', type=int, default=-1)
-parser.add_argument('--image_path', type=str, default='generated')
-parser.add_argument('--models_root', type=str, default='pretrained')
-parser.add_argument('--token_count', type=int, default=256) # for debugging
+parser.add_argument('--grid-size', type=int, default=1)
+parser.add_argument('--image-path', type=str, default='generated')
+parser.add_argument('--models-root', type=str, default='pretrained')
+parser.add_argument('--token-count', type=int, default=256) # for debugging
 
 
 def ascii_from_image(image: Image.Image, size: int) -> str:
@@ -38,6 +39,7 @@ def generate_image(
     is_mega: bool,
     text: str,
     seed: int,
+    grid_size: int,
     image_path: str,
     models_root: str,
     token_count: int
@@ -51,10 +53,10 @@ def generate_image(
     )
 
     if token_count < 256:
-        image_tokens = model.generate_image_tokens(text, seed)
-        print('image tokens', list(image_tokens.to('cpu').detach().numpy()))
+        image_tokens = model.generate_image_tokens(text, seed, grid_size ** 2)
+        print('image tokens', image_tokens.to('cpu').detach().numpy())
     else:
-        image = model.generate_image(text, seed)
+        image = model.generate_image(text, seed, grid_size)
         save_image(image, image_path)
         print(ascii_from_image(image, size=128))
 
@@ -66,6 +68,7 @@ if __name__ == '__main__':
         is_mega=args.mega,
         text=args.text,
         seed=args.seed,
+        grid_size=args.grid_size,
         image_path=args.image_path,
         models_root=args.models_root,
         token_count=args.token_count