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:
1 parent
e03082629a
commit
8d36ad5267
2 files changed
+210
-16
No files matched your search
Executable
+170
@@ -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()
|
||||
Reference in new issue
Block a user