feat: bundle gemini image generate script + update SKILL.md install guide

- Add scripts/gemini_image_generate.py to repo (was only in local skill dir)
- Update SKILL.md section 4: add step-by-step setup for image generation env
- No changes to app.py or any existing logic; script is a standalone addition
- Users can now set up image generation by following SKILL.md instructions
This commit is contained in:
Star committed 2026-03-06 01:27:30 +08:00
1 parent e03082629a
commit 8d36ad5267
2 files changed
+210 -16

No files matched your search

+170
View File
@@ -0,0 +1,170 @@
#!/usr/bin/env python3
"""Gemini Image Generate - CLI for Star Office UI background generation.
Calls Google's Gemini API to generate images, with optional reference image
for style transfer / layout preservation.
Expected interface (called by Star Office UI backend):
python gemini_image_generate.py \
--prompt "..." \
--model <model_name> \
--out-dir /tmp/xxx \
--cleanup \
[--aspect-ratio 16:9] \
[--reference-image /path/to/ref.webp]
Environment:
GEMINI_API_KEY - Google AI API key (required)
GEMINI_MODEL - override model name (optional, --model takes precedence)
Output (last line of stdout):
{"files": ["/tmp/xxx/generated_0.png"]}
"""
import argparse
import base64
import json
import mimetypes
import os
import sys
import tempfile
import shutil
from pathlib import Path
try:
from google import genai
from google.genai import types
HAS_GENAI = True
except ImportError:
HAS_GENAI = False
def detect_mime(path: str) -> str:
mt, _ = mimetypes.guess_type(path)
if mt:
return mt
ext = os.path.splitext(path)[1].lower()
return {
".png": "image/png",
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".webp": "image/webp",
".gif": "image/gif",
}.get(ext, "image/png")
def main():
parser = argparse.ArgumentParser(description="Generate image via Gemini API")
parser.add_argument("--prompt", required=True, help="Generation prompt")
parser.add_argument("--model", default="", help="Model name")
parser.add_argument("--out-dir", required=True, help="Output directory")
parser.add_argument("--cleanup", action="store_true", help="(ignored, kept for compat)")
parser.add_argument("--aspect-ratio", default="", help="Aspect ratio hint (e.g. 16:9)")
parser.add_argument("--reference-image", default="", help="Reference image path")
args = parser.parse_args()
# Resolve API key
api_key = os.environ.get("GEMINI_API_KEY", "").strip()
if not api_key:
api_key = os.environ.get("GOOGLE_API_KEY", "").strip()
if not api_key:
print("ERROR: GEMINI_API_KEY or GOOGLE_API_KEY not set", file=sys.stderr)
sys.exit(1)
# Resolve model - env var GEMINI_MODEL overrides --model flag
model = os.environ.get("GEMINI_MODEL", "").strip() or args.model.strip()
if not model:
model = "gemini-2.0-flash-exp"
# Ensure output directory
out_dir = args.out_dir
os.makedirs(out_dir, exist_ok=True)
if not HAS_GENAI:
print("ERROR: google-genai package not installed", file=sys.stderr)
sys.exit(1)
# Initialize client
client = genai.Client(api_key=api_key)
# Build prompt parts
contents = []
# Add reference image if provided
if args.reference_image and os.path.exists(args.reference_image):
ref_path = args.reference_image
mime = detect_mime(ref_path)
with open(ref_path, "rb") as f:
ref_data = f.read()
contents.append(
types.Part.from_bytes(data=ref_data, mime_type=mime)
)
# Add text prompt
prompt_text = args.prompt
if args.aspect_ratio:
prompt_text += f"\nTarget aspect ratio: {args.aspect_ratio}."
contents.append(prompt_text)
# Configure generation
generate_config = types.GenerateContentConfig(
response_modalities=["TEXT", "IMAGE"],
)
try:
response = client.models.generate_content(
model=model,
contents=contents,
config=generate_config,
)
except Exception as e:
err_msg = str(e)
print(f"ERROR: {err_msg}", file=sys.stderr)
sys.exit(1)
# Extract generated images
output_files = []
idx = 0
if response.candidates:
for candidate in response.candidates:
if not candidate.content or not candidate.content.parts:
continue
for part in candidate.content.parts:
if part.inline_data and part.inline_data.mime_type and part.inline_data.mime_type.startswith("image/"):
# Determine extension from mime
mime = part.inline_data.mime_type
ext_map = {
"image/png": ".png",
"image/jpeg": ".jpg",
"image/webp": ".webp",
}
ext = ext_map.get(mime, ".png")
out_path = os.path.join(out_dir, f"generated_{idx}{ext}")
with open(out_path, "wb") as f:
f.write(part.inline_data.data)
output_files.append(out_path)
idx += 1
if not output_files:
# Check if there's text response with error info
text_parts = []
if response.candidates:
for candidate in response.candidates:
if candidate.content and candidate.content.parts:
for part in candidate.content.parts:
if part.text:
text_parts.append(part.text)
if text_parts:
print(f"ERROR: No image generated. Model response: {' '.join(text_parts)[:500]}", file=sys.stderr)
else:
print("ERROR: No image generated and no text response", file=sys.stderr)
sys.exit(1)
# Output result as JSON (backend reads the last line)
result = {"files": output_files}
print(json.dumps(result))
if __name__ == "__main__":
main()