File size: 1,355 Bytes
e465a2f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | """
main.py - PixelModel v1 inference from model.png (the canonical model).
Usage:
python main.py "a red double decker bus"
python main.py "a cat on a couch" --out cat.png --res 64 --scale 4
"""
import argparse
import numpy as np
import torch
from PIL import Image
from model import NATIVE_RES, forward, load_model
def main():
p = argparse.ArgumentParser()
p.add_argument("prompt")
p.add_argument("--model", default="model.png")
p.add_argument("--out", default="out.png")
p.add_argument("--res", type=int, default=NATIVE_RES,
help="generation resolution (decoder is resolution-free)")
p.add_argument("--scale", type=int, default=4, help="nearest-neighbor upscale for viewing")
args = p.parse_args()
weights = load_model(args.model)
with torch.no_grad():
result = forward(weights, args.prompt, res=args.res)[0]
arr = (result.numpy() * 255).clip(0, 255).astype(np.uint8)
img = Image.fromarray(arr, mode="RGB")
if args.scale > 1:
img = img.resize((args.res * args.scale,) * 2, Image.NEAREST)
img.save(args.out)
print(f"prompt : '{args.prompt}'")
print(f"model : {args.model}")
print(f"output : {args.out} ({args.res}x{args.res} native, x{args.scale} view)")
if __name__ == "__main__":
main()
|