2026-03-10 12:05:41 +07:00

488 lines
17 KiB
Python

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Icon Generation Script using Gemini 3.1 Pro Preview API
Generates SVG icons via text generation (SVG is XML text format)
Model: gemini-3.1-pro-preview - best thinking, token efficiency, factual consistency
Usage:
python generate.py --prompt "settings gear icon" --style outlined
python generate.py --prompt "shopping cart" --style filled --color "#6366F1"
python generate.py --name "dashboard" --category navigation --style duotone
python generate.py --prompt "cloud upload" --batch 4 --output-dir ./icons
python generate.py --prompt "user profile" --sizes "16,24,32,48"
"""
import argparse
import json
import os
import re
import sys
import time
from pathlib import Path
from datetime import datetime
def load_env():
"""Load .env files in priority order"""
env_paths = [
Path(__file__).parent.parent.parent / ".env",
Path.home() / ".claude" / "skills" / ".env",
Path.home() / ".claude" / ".env"
]
for env_path in env_paths:
if env_path.exists():
with open(env_path) as f:
for line in f:
line = line.strip()
if line and not line.startswith('#') and '=' in line:
key, value = line.split('=', 1)
if key not in os.environ:
os.environ[key] = value.strip('"\'')
load_env()
try:
from google import genai
from google.genai import types
except ImportError:
print("Error: google-genai package not installed.")
print("Install with: pip install google-genai")
sys.exit(1)
# ============ CONFIGURATION ============
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
MODEL = "gemini-3.1-pro-preview"
# Icon styles with SVG-specific instructions
ICON_STYLES = {
"outlined": "outlined stroke icons, 2px stroke width, no fill, clean open paths",
"filled": "solid filled icons, no stroke, flat color fills, bold shapes",
"duotone": "duotone style with primary color at full opacity and secondary color at 30% opacity, layered shapes",
"thin": "thin line icons, 1px or 1.5px stroke width, delicate minimalist lines",
"bold": "bold thick line icons, 3px stroke width, heavy weight, impactful",
"rounded": "rounded icons with round line caps and joins, soft corners, friendly feel",
"sharp": "sharp angular icons, square line caps and mitered joins, precise edges",
"flat": "flat design icons, solid fills, no gradients or shadows, geometric simplicity",
"gradient": "linear or radial gradient fills, modern vibrant color transitions",
"glassmorphism": "glassmorphism style with semi-transparent fills, blur backdrop effect simulation, frosted glass",
"pixel": "pixel art style icons on a grid, retro 8-bit aesthetic, crisp edges",
"hand-drawn": "hand-drawn sketch style, slightly irregular strokes, organic feel, imperfect lines",
"isometric": "isometric 3D projection, 30-degree angles, dimensional depth",
"glyph": "simple glyph style, single solid shape, minimal detail, pictogram",
"animated-ready": "animated-ready SVG with named groups and IDs for CSS/JS animation targets",
}
ICON_CATEGORIES = {
"navigation": "arrows, menus, hamburger, chevrons, home, back, forward, breadcrumb",
"action": "edit, delete, save, download, upload, share, copy, paste, print, search",
"communication": "email, chat, phone, video call, notification, bell, message bubble",
"media": "play, pause, stop, skip, volume, microphone, camera, image, gallery",
"file": "document, folder, archive, attachment, cloud, database, storage",
"user": "person, group, avatar, profile, settings, lock, key, shield",
"commerce": "cart, bag, wallet, credit card, receipt, tag, gift, store",
"data": "chart, graph, analytics, dashboard, table, filter, sort, calendar",
"development": "code, terminal, bug, git, API, server, database, deploy",
"social": "heart, star, thumbs up, bookmark, flag, trophy, badge, crown",
"weather": "sun, moon, cloud, rain, snow, wind, thunder, temperature",
"map": "pin, location, compass, globe, route, directions, map marker",
}
# SVG generation prompt template
SVG_PROMPT_TEMPLATE = """Generate a clean, production-ready SVG icon.
Requirements:
- Output ONLY valid SVG code, nothing else
- ViewBox: "0 0 {viewbox} {viewbox}"
- Use currentColor for strokes/fills (inherits CSS color)
- No embedded fonts or text elements unless specifically requested
- No raster images or external references
- Optimized paths with minimal nodes
- Accessible: include <title> element with icon description
{style_instructions}
{color_instructions}
{size_instructions}
Icon to generate: {prompt}
Output the SVG code only, wrapped in ```svg``` code block."""
SVG_BATCH_PROMPT_TEMPLATE = """Generate {count} distinct SVG icon variations for: {prompt}
Requirements for EACH icon:
- Output ONLY valid SVG code
- ViewBox: "0 0 {viewbox} {viewbox}"
- Use currentColor for strokes/fills (inherits CSS color)
- No embedded fonts, raster images, or external references
- Optimized paths with minimal nodes
- Include <title> element with icon description
{style_instructions}
{color_instructions}
Generate {count} different visual interpretations. Output each SVG in a separate ```svg``` code block.
Label each variation (e.g., "Variation 1: [brief description]")."""
def extract_svgs(text):
"""Extract SVG code blocks from model response"""
svgs = []
# Try ```svg code blocks first
pattern = r'```svg\s*\n(.*?)```'
matches = re.findall(pattern, text, re.DOTALL)
if matches:
svgs.extend(matches)
# Fallback: try ```xml code blocks
if not svgs:
pattern = r'```xml\s*\n(.*?)```'
matches = re.findall(pattern, text, re.DOTALL)
svgs.extend(matches)
# Fallback: try bare <svg> tags
if not svgs:
pattern = r'(<svg[^>]*>.*?</svg>)'
matches = re.findall(pattern, text, re.DOTALL)
svgs.extend(matches)
# Clean up extracted SVGs
cleaned = []
for svg in svgs:
svg = svg.strip()
if not svg.startswith('<svg'):
# Try to find <svg> within the extracted text
match = re.search(r'(<svg[^>]*>.*?</svg>)', svg, re.DOTALL)
if match:
svg = match.group(1)
else:
continue
cleaned.append(svg)
return cleaned
def apply_color(svg_code, color):
"""Replace currentColor with specific color if provided"""
if color:
# Replace currentColor with the specified color
svg_code = svg_code.replace('currentColor', color)
# If no currentColor was present, add fill/stroke color
if color not in svg_code:
svg_code = svg_code.replace('<svg', f'<svg color="{color}"', 1)
return svg_code
def apply_viewbox_size(svg_code, size):
"""Adjust SVG viewBox to target size"""
if size:
# Update width/height attributes if present
svg_code = re.sub(r'width="[^"]*"', f'width="{size}"', svg_code)
svg_code = re.sub(r'height="[^"]*"', f'height="{size}"', svg_code)
# Add width/height if not present
if 'width=' not in svg_code:
svg_code = svg_code.replace('<svg', f'<svg width="{size}" height="{size}"', 1)
return svg_code
def generate_icon(prompt, style=None, category=None, name=None,
color=None, size=24, output_path=None, viewbox=24):
"""Generate a single SVG icon using Gemini 3.1 Pro Preview"""
if not GEMINI_API_KEY:
print("Error: GEMINI_API_KEY not set")
print("Set it with: export GEMINI_API_KEY='your-key'")
return None
client = genai.Client(api_key=GEMINI_API_KEY)
# Build style instructions
style_instructions = ""
if style and style in ICON_STYLES:
style_instructions = f"- Style: {ICON_STYLES[style]}"
# Build color instructions
color_instructions = "- Use currentColor for all strokes and fills"
if color:
color_instructions = f"- Use color: {color} for primary elements, currentColor for secondary"
# Build size instructions
size_instructions = f"- Design for {size}px display size, optimize detail level accordingly"
# Build final prompt
icon_prompt = prompt
if category and category in ICON_CATEGORIES:
icon_prompt = f"{prompt} (category: {ICON_CATEGORIES[category]})"
if name:
icon_prompt = f"'{name}' icon: {icon_prompt}"
full_prompt = SVG_PROMPT_TEMPLATE.format(
prompt=icon_prompt,
viewbox=viewbox,
style_instructions=style_instructions,
color_instructions=color_instructions,
size_instructions=size_instructions
)
print(f"Generating icon with {MODEL}...")
print(f"Prompt: {prompt}")
if style:
print(f"Style: {style}")
print()
try:
response = client.models.generate_content(
model=MODEL,
contents=full_prompt,
config=types.GenerateContentConfig(
temperature=0.7,
max_output_tokens=4096,
)
)
# Extract SVG from response
response_text = response.text if hasattr(response, 'text') else ""
if not response_text:
for part in response.candidates[0].content.parts:
if hasattr(part, 'text') and part.text:
response_text += part.text
svgs = extract_svgs(response_text)
if not svgs:
print("No valid SVG generated. Model response:")
print(response_text[:500])
return None
svg_code = svgs[0]
# Apply color if specified
svg_code = apply_color(svg_code, color)
# Apply size
svg_code = apply_viewbox_size(svg_code, size)
# Determine output path
if output_path is None:
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
slug = name or prompt.split()[0] if prompt else "icon"
slug = re.sub(r'[^a-zA-Z0-9_-]', '_', slug.lower())
style_suffix = f"_{style}" if style else ""
output_path = f"{slug}{style_suffix}_{timestamp}.svg"
# Save SVG
with open(output_path, "w", encoding="utf-8") as f:
f.write(svg_code)
print(f"Icon saved to: {output_path}")
return output_path
except Exception as e:
print(f"Error generating icon: {e}")
return None
def generate_batch(prompt, count, output_dir, style=None, color=None,
viewbox=24, name=None):
"""Generate multiple icon variations"""
if not GEMINI_API_KEY:
print("Error: GEMINI_API_KEY not set")
return []
client = genai.Client(api_key=GEMINI_API_KEY)
os.makedirs(output_dir, exist_ok=True)
# Build instructions
style_instructions = ""
if style and style in ICON_STYLES:
style_instructions = f"- Style: {ICON_STYLES[style]}"
color_instructions = "- Use currentColor for all strokes and fills"
if color:
color_instructions = f"- Use color: {color} for primary elements"
full_prompt = SVG_BATCH_PROMPT_TEMPLATE.format(
prompt=prompt,
count=count,
viewbox=viewbox,
style_instructions=style_instructions,
color_instructions=color_instructions
)
print(f"\n{'='*60}")
print(f" BATCH ICON GENERATION")
print(f" Model: {MODEL}")
print(f" Prompt: {prompt}")
print(f" Variants: {count}")
print(f" Output: {output_dir}")
print(f"{'='*60}\n")
try:
response = client.models.generate_content(
model=MODEL,
contents=full_prompt,
config=types.GenerateContentConfig(
temperature=0.9,
max_output_tokens=16384,
)
)
response_text = response.text if hasattr(response, 'text') else ""
if not response_text:
for part in response.candidates[0].content.parts:
if hasattr(part, 'text') and part.text:
response_text += part.text
svgs = extract_svgs(response_text)
if not svgs:
print("No valid SVGs generated.")
print(response_text[:500])
return []
results = []
slug = name or re.sub(r'[^a-zA-Z0-9_-]', '_', prompt.split()[0].lower())
style_suffix = f"_{style}" if style else ""
for i, svg_code in enumerate(svgs[:count]):
svg_code = apply_color(svg_code, color)
filename = f"{slug}{style_suffix}_{i+1:02d}.svg"
filepath = os.path.join(output_dir, filename)
with open(filepath, "w", encoding="utf-8") as f:
f.write(svg_code)
results.append(filepath)
print(f" [{i+1}/{len(svgs[:count])}] Saved: {filename}")
print(f"\n{'='*60}")
print(f" BATCH COMPLETE: {len(results)}/{count} icons generated")
print(f"{'='*60}\n")
return results
except Exception as e:
print(f"Error generating icons: {e}")
return []
def generate_sizes(prompt, sizes, style=None, color=None, output_dir=None, name=None):
"""Generate same icon at multiple sizes"""
if output_dir is None:
output_dir = "."
os.makedirs(output_dir, exist_ok=True)
results = []
slug = name or re.sub(r'[^a-zA-Z0-9_-]', '_', prompt.split()[0].lower())
style_suffix = f"_{style}" if style else ""
for size in sizes:
print(f"Generating {size}px variant...")
filename = f"{slug}{style_suffix}_{size}px.svg"
filepath = os.path.join(output_dir, filename)
result = generate_icon(
prompt=prompt,
style=style,
color=color,
size=size,
output_path=filepath,
viewbox=size
)
if result:
results.append(result)
time.sleep(1)
return results
def main():
parser = argparse.ArgumentParser(
description="Generate SVG icons using Gemini 3.1 Pro Preview"
)
parser.add_argument("--prompt", "-p", type=str, help="Icon description")
parser.add_argument("--name", "-n", type=str, help="Icon name (for filename)")
parser.add_argument("--style", "-s", choices=list(ICON_STYLES.keys()),
help="Icon style")
parser.add_argument("--category", "-c", choices=list(ICON_CATEGORIES.keys()),
help="Icon category for context")
parser.add_argument("--color", type=str,
help="Primary color (hex, e.g. #6366F1). Default: currentColor")
parser.add_argument("--size", type=int, default=24,
help="Icon size in px (default: 24)")
parser.add_argument("--viewbox", type=int, default=24,
help="SVG viewBox size (default: 24)")
parser.add_argument("--output", "-o", type=str, help="Output file path")
parser.add_argument("--output-dir", type=str, help="Output directory for batch")
parser.add_argument("--batch", type=int,
help="Number of icon variants to generate")
parser.add_argument("--sizes", type=str,
help="Comma-separated sizes (e.g. '16,24,32,48')")
parser.add_argument("--list-styles", action="store_true",
help="List available icon styles")
parser.add_argument("--list-categories", action="store_true",
help="List available icon categories")
args = parser.parse_args()
if args.list_styles:
print("Available icon styles:")
for style, desc in ICON_STYLES.items():
print(f" {style}: {desc[:70]}...")
return
if args.list_categories:
print("Available icon categories:")
for cat, desc in ICON_CATEGORIES.items():
print(f" {cat}: {desc}")
return
if not args.prompt and not args.name:
parser.error("Either --prompt or --name is required")
prompt = args.prompt or args.name
# Multi-size mode
if args.sizes:
sizes = [int(s.strip()) for s in args.sizes.split(",")]
generate_sizes(
prompt=prompt,
sizes=sizes,
style=args.style,
color=args.color,
output_dir=args.output_dir or "./icons",
name=args.name
)
# Batch mode
elif args.batch:
output_dir = args.output_dir or "./icons"
generate_batch(
prompt=prompt,
count=args.batch,
output_dir=output_dir,
style=args.style,
color=args.color,
viewbox=args.viewbox,
name=args.name
)
# Single icon
else:
generate_icon(
prompt=prompt,
style=args.style,
category=args.category,
name=args.name,
color=args.color,
size=args.size,
output_path=args.output,
viewbox=args.viewbox
)
if __name__ == "__main__":
main()