min-dalle-test/image_from_text.py

81 lines
2.4 KiB
Python
Raw Normal View History

2022-06-27 19:46:04 +00:00
import argparse
import os
from PIL import Image
from min_dalle import MinDalle
2022-06-27 19:46:04 +00:00
parser = argparse.ArgumentParser()
2022-06-27 20:49:42 +00:00
parser.add_argument('--mega', action='store_true')
parser.add_argument('--no-mega', dest='mega', action='store_false')
parser.set_defaults(mega=False)
2022-07-04 20:06:49 +00:00
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)
parser.add_argument('--image-path', type=str, default='generated')
parser.add_argument('--models-root', type=str, default='pretrained')
2022-07-04 13:58:00 +00:00
parser.add_argument('--row-count', type=int, default=16) # for debugging
2022-06-27 19:46:04 +00:00
2022-07-04 20:06:49 +00:00
def ascii_from_image(image: Image.Image, size: int = 128) -> str:
2022-06-27 19:46:04 +00:00
rgb_pixels = image.resize((size, int(0.55 * size))).convert('L').getdata()
chars = list('.,;/IOX')
chars = [chars[i * len(chars) // 256] for i in rgb_pixels]
chars = [chars[i * size: (i + 1) * size] for i in range(size // 2)]
return '\n'.join(''.join(row) for row in chars)
2022-06-27 20:49:42 +00:00
def save_image(image: Image.Image, path: str):
2022-06-27 19:46:04 +00:00
if os.path.isdir(path):
2022-07-04 11:38:16 +00:00
path = os.path.join(path, 'generated.jpg')
elif not path.endswith('.jpg'):
path += '.jpg'
2022-06-27 19:46:04 +00:00
print("saving image to", path)
image.save(path)
return image
def generate_image(
is_mega: bool,
text: str,
seed: int,
grid_size: int,
image_path: str,
2022-07-02 03:39:48 +00:00
models_root: str,
2022-07-04 13:58:00 +00:00
row_count: int
):
model = MinDalle(
2022-07-01 22:16:55 +00:00
is_mega=is_mega,
2022-07-02 03:39:48 +00:00
models_root=models_root,
2022-07-01 22:16:55 +00:00
is_reusable=False,
2022-07-02 00:17:20 +00:00
is_verbose=True
2022-07-01 22:16:55 +00:00
)
2022-07-04 13:58:00 +00:00
if row_count < 16:
token_count = 16 * row_count
image_tokens = model.generate_image_tokens(
text,
seed,
2022-07-04 21:27:23 +00:00
grid_size ** 2,
2022-07-04 20:06:49 +00:00
row_count,
is_verbose=True
2022-07-04 13:58:00 +00:00
)
image_tokens = image_tokens[:, :token_count].to('cpu').detach().numpy()
print('image tokens', image_tokens)
else:
2022-07-04 20:06:49 +00:00
image = model.generate_image(text, seed, grid_size, is_verbose=True)
save_image(image, image_path)
print(ascii_from_image(image, size=128))
2022-06-27 19:46:04 +00:00
if __name__ == '__main__':
args = parser.parse_args()
2022-06-27 20:49:42 +00:00
print(args)
generate_image(
is_mega=args.mega,
text=args.text,
seed=args.seed,
grid_size=args.grid_size,
image_path=args.image_path,
2022-07-02 03:39:48 +00:00
models_root=args.models_root,
2022-07-04 13:58:00 +00:00
row_count=args.row_count
)