mirror of
https://github.com/nolen777/eagle0.git
synced 2026-07-28 21:55:42 +00:00
353 lines
12 KiB
Python
353 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import base64
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import random
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
|
|
API_KEY_FILE = Path("src/main/scala/net/eagle0/common/llm_integration/api_keys.txt")
|
|
IMAGE_ENDPOINT = "https://api.openai.com/v1/images/generations"
|
|
|
|
APPEARANCE_PROFILES = [
|
|
"West African features, very dark brown skin",
|
|
"East African features, deep brown skin",
|
|
"North African features, olive-brown skin",
|
|
"Arabian Peninsula features, warm medium-brown skin",
|
|
"Persian features, olive skin",
|
|
"Turkic Central Asian features, tan skin",
|
|
"Mongolian features, golden-tan skin",
|
|
"North Indian features, medium-brown skin",
|
|
"South Indian features, dark brown skin",
|
|
"Bengali features, warm brown skin",
|
|
"Southeast Asian features, medium tan skin",
|
|
"Vietnamese features, light golden skin",
|
|
"Chinese features, light olive skin",
|
|
"Korean features, fair skin",
|
|
"Japanese features, fair-to-light olive skin",
|
|
"Indigenous Andean features, copper-brown skin",
|
|
"Mesoamerican features, medium brown skin",
|
|
"Native North American features, warm brown skin",
|
|
"Maori or Polynesian features, medium brown skin",
|
|
"Aboriginal Australian features, dark brown skin",
|
|
"Afro-Caribbean features, dark brown skin",
|
|
"Afro-Latin features, medium-dark brown skin",
|
|
"Brazilian mixed-heritage features, medium brown skin",
|
|
"Mexican mestizo features, tan skin",
|
|
"Mediterranean European features, olive skin",
|
|
"Iberian features, light olive skin",
|
|
"Italian features, olive skin",
|
|
"Greek features, olive skin",
|
|
"French features, light skin",
|
|
"Irish features, fair skin with freckles",
|
|
"Scottish features, ruddy fair skin",
|
|
"Nordic features, very fair skin",
|
|
"Slavic features, fair-to-light olive skin",
|
|
"Ashkenazi Jewish features, light olive skin",
|
|
"Armenian features, olive skin",
|
|
"Ethiopian highland features, medium-dark brown skin",
|
|
"Moroccan Amazigh features, tan olive skin",
|
|
"Egyptian features, warm tan skin",
|
|
"Filipino features, medium golden-brown skin",
|
|
"Pacific Islander features, warm brown skin",
|
|
]
|
|
|
|
AGE_PROFILES = [
|
|
"young adult",
|
|
"adult",
|
|
"seasoned middle-aged",
|
|
"older veteran",
|
|
]
|
|
|
|
FACE_PROFILES = [
|
|
"broad, sturdy face",
|
|
"narrow, severe face",
|
|
"round, watchful face",
|
|
"angular, weathered face",
|
|
"scarred but composed face",
|
|
"plain, practical face",
|
|
"commanding, heavy-browed face",
|
|
"tired but alert face",
|
|
]
|
|
|
|
HAIR_PROFILES = [
|
|
"cropped dark hair",
|
|
"close-cropped gray hair",
|
|
"braided black hair",
|
|
"shaved head",
|
|
"curly dark hair pulled back",
|
|
"straight black hair tied back",
|
|
"auburn hair streaked with gray",
|
|
"white hair cut short",
|
|
"dark beard and close hair",
|
|
"neat graying beard",
|
|
"coiled natural hair",
|
|
"long dark hair bound under a cap",
|
|
]
|
|
|
|
|
|
def api_key():
|
|
env_value = os.environ.get("OPENAI_API_KEY")
|
|
if env_value:
|
|
return env_value
|
|
|
|
if API_KEY_FILE.exists():
|
|
for line in API_KEY_FILE.read_text().splitlines():
|
|
line = line.strip()
|
|
if not line or line.startswith("#") or "=" not in line:
|
|
continue
|
|
key, value = line.split("=", 1)
|
|
if key.strip() == "openai_api_key":
|
|
value = value.strip()
|
|
if value and not value.startswith("your_"):
|
|
return value
|
|
|
|
raise RuntimeError(
|
|
"OpenAI API key not found. Set OPENAI_API_KEY or add openai_api_key to "
|
|
f"{API_KEY_FILE}."
|
|
)
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(
|
|
description="Generate Warden hero headshots and update generated_heroes.tsv image paths."
|
|
)
|
|
parser.add_argument(
|
|
"--hero-tsv",
|
|
default="src/main/resources/net/eagle0/eagle/generated_heroes.tsv",
|
|
help="Generated heroes TSV to read and update.",
|
|
)
|
|
parser.add_argument(
|
|
"--base-image-dir",
|
|
default=str(Path.home() / "Documents/headshots"),
|
|
help="Local headshot root that sync_headshots.sh uploads to S3.",
|
|
)
|
|
parser.add_argument(
|
|
"--model",
|
|
default="gpt-image-2",
|
|
help="OpenAI image generation model.",
|
|
)
|
|
parser.add_argument(
|
|
"--quality",
|
|
default="low",
|
|
choices=("low", "medium", "high", "auto"),
|
|
help="Image generation quality.",
|
|
)
|
|
parser.add_argument(
|
|
"--size",
|
|
default="1024x1024",
|
|
help="Generated image size.",
|
|
)
|
|
parser.add_argument(
|
|
"--limit",
|
|
type=int,
|
|
default=0,
|
|
help="Generate at most this many missing images; 0 means all missing images.",
|
|
)
|
|
parser.add_argument(
|
|
"--dry-run",
|
|
action="store_true",
|
|
help="Print planned paths without calling the API or writing files.",
|
|
)
|
|
parser.add_argument(
|
|
"--overwrite",
|
|
action="store_true",
|
|
help="Regenerate images even if the target PNG already exists.",
|
|
)
|
|
parser.add_argument(
|
|
"--concurrency",
|
|
type=int,
|
|
default=4,
|
|
help="Number of parallel image generation requests.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def read_rows(path):
|
|
lines = path.read_text().splitlines()
|
|
fieldnames = lines[0].split("\t")
|
|
rows = []
|
|
for line in lines[1:]:
|
|
fields = line.split("\t")
|
|
row = dict(zip(fieldnames, fields))
|
|
row["_fields"] = fields
|
|
rows.append(row)
|
|
return fieldnames, rows
|
|
|
|
|
|
def write_rows(path, fieldnames, rows):
|
|
image_path_index = fieldnames.index("image_path")
|
|
lines = ["\t".join(fieldnames)]
|
|
for row in rows:
|
|
fields = list(row["_fields"])
|
|
fields[image_path_index] = row["image_path"]
|
|
lines.append("\t".join(fields))
|
|
path.write_text("\n".join(lines) + "\n")
|
|
|
|
|
|
def image_gender(row, warden_index):
|
|
gender = row["gender"]
|
|
if gender in ("male", "female"):
|
|
return gender
|
|
|
|
digest = hashlib.sha256(row["name"].encode("utf-8")).digest()
|
|
return "male" if (digest[0] + warden_index) % 2 == 0 else "female"
|
|
|
|
|
|
def image_filename(row, warden_index):
|
|
digest = hashlib.sha256(row["name"].encode("utf-8")).hexdigest()[:16]
|
|
return f"{warden_index:04d}_{digest}.png"
|
|
|
|
|
|
def choose_profile(items, row, warden_index, salt):
|
|
digest = hashlib.sha256(f"{salt}:{warden_index}:{row['name']}".encode("utf-8")).digest()
|
|
return items[int.from_bytes(digest[:4], "big") % len(items)]
|
|
|
|
|
|
def prompt_for(row, warden_index):
|
|
prime_notes = [
|
|
f"strength {row['strength']}",
|
|
f"agility {row['agility']}",
|
|
f"constitution {row['constitution']}",
|
|
f"wisdom {row['wisdom']}",
|
|
f"charisma {row['charisma']}",
|
|
]
|
|
appearance = choose_profile(APPEARANCE_PROFILES, row, warden_index, "appearance")
|
|
age = choose_profile(AGE_PROFILES, row, warden_index, "age")
|
|
face = choose_profile(FACE_PROFILES, row, warden_index, "face")
|
|
hair = choose_profile(HAIR_PROFILES, row, warden_index, "hair")
|
|
return (
|
|
"Fantasy strategy game hero portrait, square headshot composition. "
|
|
f"Create a distinctive medieval Warden named {row['name']}. "
|
|
f"Gender presentation: {row['gender']}. "
|
|
f"Appearance: {age}, {appearance}, {face}, {hair}. "
|
|
f"Personality traits: {row['personality']}. "
|
|
"The character is a jailer, keeper of prisoners, and battlefield custodian: "
|
|
"stern keys, austere authority, prison-watch uniform, no modern clothing. "
|
|
"Show a resilient, high-constitution figure with weathered endurance and moral judgment. "
|
|
f"Stat cues: {', '.join(prime_notes)}. "
|
|
"Painted realistic fantasy style, expressive face, shoulders visible, neutral background, "
|
|
"soft dramatic lighting, no text, no logo, no frame, no watermark."
|
|
)
|
|
|
|
|
|
def request_image(key, model, quality, size, prompt):
|
|
body = {
|
|
"model": model,
|
|
"prompt": prompt,
|
|
"size": size,
|
|
"quality": quality,
|
|
"n": 1,
|
|
"output_format": "png",
|
|
}
|
|
request = urllib.request.Request(
|
|
IMAGE_ENDPOINT,
|
|
data=json.dumps(body).encode("utf-8"),
|
|
headers={
|
|
"Authorization": f"Bearer {key}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(request, timeout=300) as response:
|
|
response_body = json.loads(response.read().decode("utf-8"))
|
|
|
|
data = response_body["data"][0]
|
|
if "b64_json" in data:
|
|
return base64.b64decode(data["b64_json"])
|
|
|
|
if "url" in data:
|
|
with urllib.request.urlopen(data["url"], timeout=300) as response:
|
|
return response.read()
|
|
|
|
raise RuntimeError("Image response did not include b64_json or url")
|
|
|
|
|
|
def generate_with_retries(key, args, prompt, attempts=5):
|
|
for attempt in range(1, attempts + 1):
|
|
try:
|
|
return request_image(key, args.model, args.quality, args.size, prompt)
|
|
except urllib.error.HTTPError as error:
|
|
body = error.read().decode("utf-8", errors="replace")
|
|
retryable = error.code in (408, 409, 429, 500, 502, 503, 504)
|
|
if not retryable or attempt == attempts:
|
|
raise RuntimeError(f"OpenAI image request failed: {error.code} {body}") from error
|
|
except (TimeoutError, urllib.error.URLError):
|
|
if attempt == attempts:
|
|
raise
|
|
|
|
time.sleep(min(60, 2**attempt) + random.random())
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
hero_tsv = Path(args.hero_tsv)
|
|
base_image_dir = Path(args.base_image_dir).expanduser()
|
|
fieldnames, rows = read_rows(hero_tsv)
|
|
|
|
key = None if args.dry_run else api_key()
|
|
warden_index = 0
|
|
generation_tasks = []
|
|
|
|
for row in rows:
|
|
if row["profession"] != "warden":
|
|
continue
|
|
|
|
warden_index += 1
|
|
gender = image_gender(row, warden_index)
|
|
rel_path = f"warden/{gender}/{image_filename(row, warden_index)}"
|
|
target_path = base_image_dir / rel_path
|
|
row["image_path"] = rel_path
|
|
|
|
if target_path.exists() and not args.overwrite:
|
|
print(f"exists {rel_path}", flush=True)
|
|
continue
|
|
|
|
generation_tasks.append(
|
|
(row["name"], rel_path, target_path, prompt_for(row, warden_index))
|
|
)
|
|
|
|
if args.dry_run:
|
|
print(f"Dry run planned {warden_index} Warden image paths", flush=True)
|
|
else:
|
|
if args.limit:
|
|
generation_tasks = generation_tasks[: args.limit]
|
|
|
|
for _, _, target_path, _ in generation_tasks:
|
|
target_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
def generate_task(task):
|
|
name, rel_path, target_path, prompt = task
|
|
print(f"generate {rel_path} for {name}", flush=True)
|
|
image_bytes = generate_with_retries(key, args, prompt)
|
|
if not image_bytes.startswith(b"\x89PNG\r\n\x1a\n"):
|
|
raise RuntimeError(f"Generated image for {name} is not a PNG")
|
|
target_path.write_bytes(image_bytes)
|
|
return rel_path
|
|
|
|
generated_count = 0
|
|
if generation_tasks:
|
|
with ThreadPoolExecutor(max_workers=max(1, args.concurrency)) as executor:
|
|
futures = [executor.submit(generate_task, task) for task in generation_tasks]
|
|
for future in as_completed(futures):
|
|
rel_path = future.result()
|
|
generated_count += 1
|
|
print(
|
|
f"wrote {rel_path} ({generated_count}/{len(generation_tasks)})",
|
|
flush=True,
|
|
)
|
|
|
|
write_rows(hero_tsv, fieldnames, rows)
|
|
print(f"Updated {hero_tsv}; generated {generated_count} missing images", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|