55 lines
2.2 KiB
Python
55 lines
2.2 KiB
Python
#!/usr/bin/env python3
|
|
# Generate bestiary portraits with fal-ai/nano-banana-pro/edit; the reference image is the
|
|
# front and front three-quarter view of the creature's 64 px idle sheet, scaled up 6x.
|
|
import base64, io, json, os, sys, urllib.request, concurrent.futures as cf
|
|
from PIL import Image
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from portraits import CREATURES, prompt_for, sheet
|
|
|
|
OUT = "/home/zink/Projects/dungeoncraft/art/generated/portraits"
|
|
KEY = open(os.path.expanduser("~/.config/dungeoncraft/fal-key")).read().strip()
|
|
EXTRA = os.environ.get("EXTRA", "")
|
|
|
|
|
|
def ref_uri(cid):
|
|
im = Image.open(sheet(cid)).convert("RGBA")
|
|
h = im.height
|
|
two = im.crop((0, 0, 2 * h, h)).resize((12 * h, 6 * h), Image.NEAREST)
|
|
bg = Image.new("RGBA", two.size, (40, 36, 44, 255))
|
|
bg.alpha_composite(two)
|
|
buf = io.BytesIO()
|
|
bg.convert("RGB").save(buf, "PNG")
|
|
return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
|
|
|
|
|
|
def gen(cid, suffix=""):
|
|
body = {"prompt": prompt_for(cid, EXTRA), "image_urls": [ref_uri(cid)], "aspect_ratio": "1:1",
|
|
"resolution": "1K", "num_images": 1, "output_format": "png"}
|
|
req = urllib.request.Request("https://fal.run/fal-ai/nano-banana-pro/edit", data=json.dumps(body).encode(),
|
|
headers={"Authorization": "Key " + KEY, "Content-Type": "application/json"})
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=600) as r:
|
|
res = json.load(r)
|
|
with urllib.request.urlopen(res["images"][0]["url"], timeout=300) as r:
|
|
data = r.read()
|
|
path = f"{OUT}/{cid}-raw{suffix}.png"
|
|
open(path, "wb").write(data)
|
|
return cid, "ok " + path
|
|
except Exception as e:
|
|
msg = str(e)
|
|
if hasattr(e, "read"):
|
|
try:
|
|
msg += " " + e.read().decode()[:300]
|
|
except Exception:
|
|
pass
|
|
return cid, "ERROR " + msg
|
|
|
|
|
|
if __name__ == "__main__":
|
|
ids = sys.argv[1:] or [c["id"] for c in CREATURES]
|
|
suffix = os.environ.get("SUFFIX", "")
|
|
with cf.ThreadPoolExecutor(4) as ex:
|
|
for cid, st in ex.map(lambda i: gen(i, suffix), ids):
|
|
print(cid, st, flush=True)
|