storytellurx/scripts/image-generation/nanogpt-quick.py

5 lines
2 KiB
Python

import base64, os, httpx; from pathlib import Path; from datetime import datetime; from dotenv import load_dotenv; from InquirerPy import inquirer
root = Path(__file__).resolve().parents[2]; load_dotenv(root / ".env"); headers = {"Authorization": f"Bearer {os.environ['NANOGPT_API_KEY']}"}; api = "https://api.nano-gpt.com"
models = httpx.get(f"{api}/api/v1/images/models", headers=headers, timeout=30).raise_for_status().json()["data"]; choices = [{"name": f"{m.get('name', m['id'])} · ${rate[1]:.6f}/image · {rate[0]}", "value": (m["id"], rate[0], rate[1])} for m in models if m.get("capabilities", {}).get("image_generation") and (prices := m.get("pricing", {}).get("per_image")) and (rate := min(prices.items(), key=lambda item: item[1]))[1] > 0]; model, size, estimate = inquirer.fuzzy(message="Search image model", choices=choices).execute()
prompt = inquirer.text(message="Prompt").execute(); negative = inquirer.text(message="Negative prompt (optional)").execute(); steps = int(inquirer.text(message="Steps", default="4").execute()); cfg = float(inquirer.text(message="CFG scale", default="2").execute())
response = httpx.post(f"{api}/v1/images/generations", headers=headers, json={"model": model, "prompt": prompt, "negative_prompt": negative, "num_inference_steps": steps, "guidance_scale": cfg, "size": size, "n": 1, "response_format": "b64_json"}, timeout=180).raise_for_status().json(); image = response["data"][0]; blob = base64.b64decode(image["b64_json"]) if image.get("b64_json") else httpx.get(image["url"], timeout=60).raise_for_status().content; ext = ".png" if blob.startswith(b"\x89PNG") else ".webp" if blob.startswith(b"RIFF") else ".jpg"; path = root / ".cache" / "nanogpt-images" / f"{datetime.now():%Y%m%d-%H%M%S}-{model.replace('/', '-')}{ext}"; path.parent.mkdir(parents=True, exist_ok=True); path.write_bytes(blob); print(f"Saved: {path}\nModel: {model} · {size} · {steps} steps · CFG {cfg}\nListed: ${estimate:.6f} · Charged: {response.get('cost', 'not reported')} USD")