prefetch_cache.py

"""Pre-generate the 12 example combos (2 faces × 6 clothes) and save PNGs into static/cached/.
Idempotent: skips files that already exist. Run once, commit the results.
"""
from __future__ import annotations

import asyncio
import os
import sys
from pathlib import Path

import httpx
from dotenv import load_dotenv

HERE = Path(__file__).resolve().parent
STATIC_DIR = HERE / "static"
FACE_DIR = STATIC_DIR / "face"
CLOTHES_DIR = STATIC_DIR / "clothes"
CACHE_DIR = STATIC_DIR / "cached"
CACHE_DIR.mkdir(exist_ok=True)

# Make sibling modules importable
sys.path.insert(0, str(HERE))
from fashn import FashnClient, FashnError, to_data_url  # noqa: E402

load_dotenv(HERE / ".env")
API_KEY = os.environ.get("FASHN_API_KEY")
if not API_KEY:
    raise SystemExit("FASHN_API_KEY missing — make sure code/.env is set")

FACES = ["tony", "dahlia"]
CLOTHES = ["1", "2", "3", "4", "5", "6"]


async def fetch_to_file(url: str, dest: Path) -> None:
    async with httpx.AsyncClient(timeout=60.0) as c:
        r = await c.get(url)
    r.raise_for_status()
    dest.write_bytes(r.content)


async def main() -> None:
    client = FashnClient(api_key=API_KEY)

    todo: list[tuple[str, str]] = []
    for f in FACES:
        for cl in CLOTHES:
            cached = CACHE_DIR / f"{f}_{cl}.png"
            if cached.exists():
                print(f"[skip] {cached.name} already exists")
                continue
            todo.append((f, cl))

    if not todo:
        print("\nAll 12 combos cached. Nothing to do.")
        return

    print(f"\nGenerating {len(todo)} combos…\n")
    for i, (face, clothing) in enumerate(todo, 1):
        cached = CACHE_DIR / f"{face}_{clothing}.png"
        face_path = FACE_DIR / f"{face}.jpg"
        clothing_path = CLOTHES_DIR / f"{clothing}.webp"

        print(f"[{i}/{len(todo)}] face={face} clothing={clothing} …", flush=True)
        try:
            result = await client.product_to_model(
                product_image=to_data_url(clothing_path),
                face_reference=to_data_url(face_path),
                aspect_ratio="3:4",
                generation_mode="balanced",
                resolution="1k",
                num_images=1,
                output_format="png",
            )
        except FashnError as e:
            print(f"   FAILED: {e}")
            continue

        output = result.get("output") or []
        if not output:
            print("   FAILED: no output returned")
            continue

        try:
            await fetch_to_file(output[0], cached)
            print(f"   ✓ saved {cached.name}")
        except Exception as e:
            print(f"   FAILED to download: {e}")

    print("\nDone.")


if __name__ == "__main__":
    asyncio.run(main())