File size: 1,724 Bytes
38bfe91 | 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 45 46 47 48 49 50 51 52 53 54 | """
main.py – run the pixel model.
Usage:
python main.py "a red circle"
python main.py "a red circle" --model model.png --out out.png --scale 8
"""
import argparse
import sys
import os
import numpy as np
from PIL import Image
import torch
from model import load_model, forward, OUT_SIZE
def generate(prompt: str, model_path: str, out_path: str, scale: int = 8):
if not os.path.exists(model_path):
sys.exit(f"Model not found: {model_path}\n"
f"Run: python train.py to create one first.")
pixels = load_model(model_path) # (H, W, 3)
with torch.no_grad():
result = forward(pixels, prompt) # (32, 32, 3)
arr = (result.numpy() * 255).clip(0, 255).astype(np.uint8)
img = Image.fromarray(arr, mode="RGB")
# upscale for visibility (nearest-neighbour keeps pixel look)
if scale > 1:
img = img.resize(
(OUT_SIZE * scale, OUT_SIZE * scale),
Image.NEAREST
)
img.save(out_path)
print(f"prompt : '{prompt}'")
print(f"model : {model_path} ({os.path.getsize(model_path)//1024} KB)")
print(f"output : {out_path} ({OUT_SIZE*scale}×{OUT_SIZE*scale} px)")
if __name__ == "__main__":
p = argparse.ArgumentParser(description="PixelModel inference")
p.add_argument("prompt", help="Text prompt")
p.add_argument("--model", default="model.png", help="Path to model PNG")
p.add_argument("--out", default="out.png", help="Output image path")
p.add_argument("--scale", type=int, default=8,
help="Upscale factor for output (default 8 → 256×256)")
args = p.parse_args()
generate(args.prompt, args.model, args.out, args.scale)
|