import argparse
from pathlib import Path

from pipeline import SpriteFlowPipeline, make_grid


def main():
    p = argparse.ArgumentParser(description="Generate pixel-art sprites from text.")
    p.add_argument("--prompt", required=True, help="Text prompt")
    p.add_argument("--model", default=str(Path(__file__).parent),
                   help="Local folder or Hugging Face repo id")
    p.add_argument("--num_images", type=int, default=8)
    p.add_argument("--steps", type=int, default=50)
    p.add_argument("--guidance", type=float, default=3.0)
    p.add_argument("--seed", type=int, default=None)
    p.add_argument("--size", type=int, default=128, help="Output size in pixels (nearest-neighbour upscale)")
    p.add_argument("--out", default="sprites.png")
    args = p.parse_args()

    pipe = SpriteFlowPipeline.from_pretrained(args.model)
    images = pipe(args.prompt, num_images=args.num_images, num_steps=args.steps,
                  guidance_scale=args.guidance, seed=args.seed, output_size=args.size)
    make_grid(images, cols=min(4, len(images))).save(args.out)
    print("saved", args.out)


if __name__ == "__main__":
    main()