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

add fp16 argument

Browse Source
This commit is contained in:
Jonas Konschak 2022-07-14 20:22:06 +02:00
parent cc78846506
commit 71fb0c29bc
1 changed files with 8 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
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)
@ -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,
) )