Files
JARVIS/gen_face_ai.py
David Rice a755e12ea3 Changes
2026-04-16 15:46:56 +01:00

92 lines
3.6 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env python3
"""
gen_face_ai.py Generate assets/face_wire.png via Gemini image generation
──────────────────────────────────────────────────────────────────────────────
Run once to produce the face PNG that display.py loads at startup.
pip install google-genai pillow
python3 gen_face_ai.py
Pass --list-models to see what image-capable models your key can reach.
"""
import os
import sys
import io
import argparse
# ── API key ───────────────────────────────────────────────────────────────────
API_KEY = 'AQ.Ab8RN6LuGwkGiKPa61jsLAEYEpJp1Yl2EkZuBWTbN9AMKxgTSw'
# ─────────────────────────────────────────────────────────────────────────────
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')
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')
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.')
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 = (
"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"
)
print(f'Generating with {args.model}')
response = client.models.generate_content(
model = args.model,
contents = PROMPT,
config = types.GenerateContentConfig(
response_modalities = ['IMAGE', 'TEXT'],
),
)
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:
for part in response.candidates[0].content.parts:
if hasattr(part, 'text') and part.text:
print('Model said:', part.text[:400])
sys.exit('No image in response try a different --model')
os.makedirs(os.path.dirname(args.out) if os.path.dirname(args.out) else '.', exist_ok=True)
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.')