min-dalle-test/image_from_text.py

54 lines
1.7 KiB
Python
Raw Normal View History

2022-06-27 19:46:04 +00:00
import argparse
import os
from PIL import Image
2022-06-27 20:49:42 +00:00
from min_dalle.generate_image import generate_image_from_text
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)
parser.add_argument('--torch', action='store_true')
parser.add_argument('--no-torch', dest='torch', action='store_false')
parser.set_defaults(torch=False)
parser.add_argument('--text', type=str)
parser.add_argument('--seed', type=int, default=0)
parser.add_argument('--image_path', type=str, default='generated')
parser.add_argument('--image_token_count', type=int, default=256) # for debugging
2022-06-27 19:46:04 +00:00
def ascii_from_image(image: Image.Image, size: int) -> str:
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):
path = os.path.join(path, 'generated.png')
elif not path.endswith('.png'):
path += '.png'
print("saving image to", path)
image.save(path)
return image
if __name__ == '__main__':
args = parser.parse_args()
2022-06-27 20:49:42 +00:00
print(args)
image = generate_image_from_text(
text = args.text,
is_mega = args.mega,
is_torch = args.torch,
seed = args.seed,
image_token_count = args.image_token_count
)
if image != None:
save_image(image, args.image_path)
2022-06-27 19:46:04 +00:00
print(ascii_from_image(image, size=128))