107 lines
5.3 KiB
Python
107 lines
5.3 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
gen_face_ai.py – Generate assets/face_wire.png using Gemini image generation
|
||
────────────────────────────────────────────────────────────────────────────────
|
||
Uses the Gemini / Google Gen AI SDK with a native image-generation model
|
||
(gemini-2.5-flash-image or similar) to produce a wireframe face PNG ready for
|
||
display.py to load.
|
||
|
||
pip install google-genai pillow
|
||
python3 gen_face_ai.py
|
||
|
||
Or set the key in the environment:
|
||
export GEMINI_API_KEY=YOUR_KEY
|
||
python3 gen_face_ai.py
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import io
|
||
import argparse
|
||
|
||
# ── Paste your Gemini API key here ────────────────────────────────────────────
|
||
API_KEY = 'AQ.Ab8RN6LuGwkGiKPa61jsLAEYEpJp1Yl2EkZuBWTbN9AMKxgTSw'
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
|
||
# ── CLI ───────────────────────────────────────────────────────────────────────
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument('--key', default='', help='Override the hardcoded API key')
|
||
ap.add_argument('--out', default='assets/face_wire.png')
|
||
ap.add_argument('--model', default='gemini-2.5-flash-image',
|
||
help='Gemini image model to use (default: gemini-2.5-flash-image)')
|
||
ap.add_argument('--list-models', action='store_true',
|
||
help='Print models that support generateContent then exit')
|
||
args = ap.parse_args()
|
||
|
||
api_key = args.key or API_KEY or os.environ.get('GEMINI_API_KEY', '')
|
||
if not api_key:
|
||
sys.exit('ERROR: paste your key into API_KEY at the top of this file')
|
||
|
||
# ── Install check ─────────────────────────────────────────────────────────────
|
||
try:
|
||
from google import genai
|
||
from google.genai import types
|
||
except ImportError:
|
||
sys.exit('Run: pip install google-genai then try again.')
|
||
|
||
try:
|
||
from PIL import Image
|
||
except ImportError:
|
||
sys.exit('Run: pip install pillow then try again.')
|
||
|
||
# ── Connect ───────────────────────────────────────────────────────────────────
|
||
print('Connecting to Google GenAI …')
|
||
client = genai.Client(api_key=api_key)
|
||
|
||
if args.list_models:
|
||
print('\nModels with generateContent support:')
|
||
for m in client.models.list():
|
||
methods = getattr(m, 'supported_actions', None) or getattr(m, 'supported_methods', None) or []
|
||
if 'generateContent' in methods:
|
||
print(f' {m.name}')
|
||
sys.exit(0)
|
||
|
||
# ── Prompt ────────────────────────────────────────────────────────────────────
|
||
PROMPT = (
|
||
"3D wireframe polygon mesh of a human head and face, viewed from slightly "
|
||
"below, front-facing, neutral expression, pure black background, thin "
|
||
"light gray lines only forming a grid over the head and face, no colour "
|
||
"fill, no skin texture, no neck, symmetrical, sci-fi holographic display "
|
||
"style, high contrast monochrome"
|
||
)
|
||
|
||
# ── Generate ──────────────────────────────────────────────────────────────────
|
||
print(f'Generating with {args.model} …')
|
||
response = client.models.generate_content(
|
||
model = args.model,
|
||
contents= PROMPT,
|
||
config = types.GenerateContentConfig(
|
||
response_modalities = ['IMAGE', 'TEXT'],
|
||
),
|
||
)
|
||
|
||
# ── Extract image bytes ───────────────────────────────────────────────────────
|
||
img_bytes = None
|
||
for part in response.candidates[0].content.parts:
|
||
if part.inline_data and part.inline_data.mime_type.startswith('image/'):
|
||
img_bytes = part.inline_data.data
|
||
break
|
||
|
||
if img_bytes is None:
|
||
# Print any text the model returned to help debug
|
||
for part in response.candidates[0].content.parts:
|
||
if hasattr(part, 'text') and part.text:
|
||
print('Model text response:', part.text[:400])
|
||
sys.exit('No image in response – try a different --model')
|
||
|
||
# ── Save as PNG ───────────────────────────────────────────────────────────────
|
||
os.makedirs(os.path.dirname(args.out) if os.path.dirname(args.out) else '.',
|
||
exist_ok=True)
|
||
|
||
# Convert whatever format came back to a proper PNG
|
||
img = Image.open(io.BytesIO(img_bytes))
|
||
img.save(args.out, 'PNG')
|
||
|
||
print(f'Saved {args.out} ({img.width}×{img.height} px)')
|
||
print('Run python3 display.py to see it.')
|