From 71fb0c29bcb52f18c11ee056baf7d3943077f0c1 Mon Sep 17 00:00:00 2001
From: Jonas Konschak <jonaskonschak@gmx.de>
Date: Thu, 14 Jul 2022 20:22:06 +0200
Subject: [PATCH 1/2] add fp16 argument

---
 image_from_text.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/image_from_text.py b/image_from_text.py
index 2d1df22..87c8a89 100644
--- a/image_from_text.py
+++ b/image_from_text.py
@@ -2,12 +2,13 @@ import argparse
 import os
 from PIL import Image
 from min_dalle import MinDalle
-
+import torch
 
 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('--fp16', action='store_true')
 parser.add_argument('--text', type=str, default='Dali painting of WALL·E')
 parser.add_argument('--seed', type=int, default=-1)
 parser.add_argument('--grid-size', type=int, default=1)
@@ -41,13 +42,15 @@ def generate_image(
     grid_size: int,
     top_k: int,
     image_path: str,
-    models_root: str
+    models_root: str,
+    fp16: bool,
 ):
     model = MinDalle(
         is_mega=is_mega, 
         models_root=models_root,
         is_reusable=False,
-        is_verbose=True
+        is_verbose=True,
+        dtype=torch.float16 if fp16 else torch.float32
     )
 
     image = model.generate_image(
@@ -71,5 +74,6 @@ if __name__ == '__main__':
         grid_size=args.grid_size,
         top_k=args.top_k,
         image_path=args.image_path,
-        models_root=args.models_root
+        models_root=args.models_root,
+        fp16=args.fp16,
     )
\ No newline at end of file

From 1ab70c790bf9e3d162e02047e49fa3116270f546 Mon Sep 17 00:00:00 2001
From: Jonas Konschak <90077736+kanttouchthis@users.noreply.github.com>
Date: Fri, 15 Jul 2022 01:24:04 +0200
Subject: [PATCH 2/2] save as png

---
 image_from_text.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/image_from_text.py b/image_from_text.py
index 87c8a89..06cf691 100644
--- a/image_from_text.py
+++ b/image_from_text.py
@@ -27,9 +27,9 @@ def ascii_from_image(image: Image.Image, size: int = 128) -> str:
 
 def save_image(image: Image.Image, path: str):
     if os.path.isdir(path):
-        path = os.path.join(path, 'generated.jpg')
-    elif not path.endswith('.jpg'):
-        path += '.jpg'
+        path = os.path.join(path, 'generated.png')
+    elif not path.endswith('.png'):
+        path += '.png'
     print("saving image to", path)
     image.save(path)
     return image
@@ -76,4 +76,4 @@ if __name__ == '__main__':
         image_path=args.image_path,
         models_root=args.models_root,
         fp16=args.fp16,
-    )
\ No newline at end of file
+    )