try_it.py

"""Quick smoke test — runs the full aipixel pipeline with real API calls.

Usage:
    set REPLICATE_API_TOKEN=r8_your_token_here
    python try_it.py
"""
import os
from pathlib import Path

from aipixel import AiPixel, GenerationConfig

TOKEN = os.environ.get("REPLICATE_API_TOKEN")
if not TOKEN:
    raise SystemExit("Set REPLICATE_API_TOKEN first:\n  set REPLICATE_API_TOKEN=r8_...")

OUT = Path("output")
OUT.mkdir(exist_ok=True)

ap = AiPixel(api_token=TOKEN)

REF = Path("reference/environment/rocks/rock_001.png")
config = GenerationConfig(width=64, height=64, style="low_res", remove_bg=True)

# ── 1. Text prompt only ────────────────────────────────────────────────────────
print("1/4  Generating from text prompt only...")
imgs = ap.generate("dark mossy rock, top-down pixel art", config=config)
imgs[0].save(OUT / "01_text_only.png")
print(f"     Saved → {OUT/'01_text_only.png'}  ({imgs[0].size} {imgs[0].mode})")

# ── 2. Text + reference image ──────────────────────────────────────────────────
print("2/4  Generating with reference image...")
imgs = ap.generate("dark rock formation", reference=REF, config=config)
imgs[0].save(OUT / "02_with_reference.png")
print(f"     Saved → {OUT/'02_with_reference.png'}  ({imgs[0].size} {imgs[0].mode})")

# ── 3. Extracted palette + reference ──────────────────────────────────────────
print("3/4  Extracting palette from reference, then generating...")
palette = ap.extract_palette(REF, max_colors=6)
print(f"     Palette: {palette}")
imgs = ap.generate("dark rock pile", reference=REF, palette=palette, config=config)
imgs[0].save(OUT / "03_with_palette.png")
print(f"     Saved → {OUT/'03_with_palette.png'}  ({imgs[0].size} {imgs[0].mode})")

# ── 4. Variations ──────────────────────────────────────────────────────────────
print("4/4  Generating 2 variations of the reference...")
variations = ap.variations(REF, count=2, prompt="dark rock, pixel art")
for i, img in enumerate(variations):
    path = OUT / f"04_variation_{i+1}.png"
    img.save(path)
    print(f"     Saved → {path}  ({img.size} {img.mode})")

print(f"\nDone. Check the '{OUT}/' folder for results.")