main.py

"""
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)