moved flax model and conversion code to separate repository

image_from_text.py

@@ -3,15 +3,11 @@ import os
 from PIL import Image
 
 from min_dalle.min_dalle_torch import MinDalleTorch
-from min_dalle.min_dalle_flax import MinDalleFlax
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--mega', action='store_true')
 parser.add_argument('--no-mega', dest='mega', action='store_false')
 parser.set_defaults(mega=False)
-parser.add_argument('--torch', action='store_true')
-parser.add_argument('--no-torch', dest='torch', action='store_false')
-parser.set_defaults(torch=True)
 parser.add_argument('--text', type=str, default='cat')
 parser.add_argument('--seed', type=int, default=0)
 parser.add_argument('--image_path', type=str, default='generated')
@@ -37,7 +33,6 @@ def save_image(image: Image.Image, path: str):
 
 
 def generate_image(
-    is_torch: bool,
     is_mega: bool,
     text: str,
     seed: int,
@@ -45,29 +40,21 @@ def generate_image(
     token_count: int
 ):
     is_reusable = False
-    if is_torch:
-        image_generator = MinDalleTorch(is_mega, is_reusable, token_count)
-
-        if token_count < 256:
-            image_tokens = image_generator.generate_image_tokens(text, seed)
-            print('image tokens', list(image_tokens.to('cpu').detach().numpy()))
-            return
-        else:
-            image = image_generator.generate_image(text, seed)
+    model = MinDalleTorch(is_mega, is_reusable, token_count)
 
+    if token_count < 256:
+        image_tokens = model.generate_image_tokens(text, seed)
+        print('image tokens', list(image_tokens.to('cpu').detach().numpy()))
     else:
-        image_generator = MinDalleFlax(is_mega, is_reusable)
-        image = image_generator.generate_image(text, seed)
-
-    save_image(image, image_path)
-    print(ascii_from_image(image, size=128))
+        image = model.generate_image(text, seed)
+        save_image(image, image_path)
+        print(ascii_from_image(image, size=128))
 
 
 if __name__ == '__main__':
     args = parser.parse_args()
     print(args)
     generate_image(
-        is_torch=args.torch,
         is_mega=args.mega,
         text=args.text,
         seed=args.seed,