tanner/min-dalle-test
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Merge pull request #75 from kanttouchthis/main

Browse Source 
add fp16 argument
This commit is contained in:
Brett Kuprel 2022-07-15 08:53:59 -04:00 committed by GitHub
commit ddfc806dd5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 12 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
image_from_text.py
View File

@ -2,12 +2,13 @@ import argparse
import os import os
from PIL import Image from PIL import Image
from min_dalle import MinDalle from min_dalle import MinDalle
import torch
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('--mega', action='store_true') parser.add_argument('--mega', action='store_true')
parser.add_argument('--no-mega', dest='mega', action='store_false') parser.add_argument('--no-mega', dest='mega', action='store_false')
parser.set_defaults(mega=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('--text', type=str, default='Dali painting of WALL·E')
parser.add_argument('--seed', type=int, default=-1) parser.add_argument('--seed', type=int, default=-1)
parser.add_argument('--grid-size', type=int, default=1) parser.add_argument('--grid-size', type=int, default=1)
@ -26,9 +27,9 @@ def ascii_from_image(image: Image.Image, size: int = 128) -> str:
def save_image(image: Image.Image, path: str): def save_image(image: Image.Image, path: str):
if os.path.isdir(path): if os.path.isdir(path):
path = os.path.join(path, 'generated.jpg') path = os.path.join(path, 'generated.png')
elif not path.endswith('.jpg'): elif not path.endswith('.png'):
path += '.jpg' path += '.png'
print("saving image to", path) print("saving image to", path)
image.save(path) image.save(path)
return image return image
@ -41,13 +42,15 @@ def generate_image(
grid_size: int, grid_size: int,
top_k: int, top_k: int,
image_path: str, image_path: str,
models_root: str models_root: str,
fp16: bool,
): ):
model = MinDalle( model = MinDalle(
is_mega=is_mega, is_mega=is_mega,
models_root=models_root, models_root=models_root,
is_reusable=False, is_reusable=False,
is_verbose=True is_verbose=True,
dtype=torch.float16 if fp16 else torch.float32
) )
image = model.generate_image( image = model.generate_image(
@ -71,5 +74,6 @@ if __name__ == '__main__':
grid_size=args.grid_size, grid_size=args.grid_size,
top_k=args.top_k, top_k=args.top_k,
image_path=args.image_path, image_path=args.image_path,
models_root=args.models_root models_root=args.models_root,
) fp16=args.fp16,
)