#!/usr/bin/env python3
"""通过 API易 调用 Nano Banana 2（gemini-3.1-flash-image-preview）生成 / 编辑图片。纯标准库，零依赖。"""
import argparse
import base64
import json
import os
import sys
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor
# 一次调用最多并发出几张图（边界，避免一次性打太多请求）
MAX_COUNT = 5
def load_api_key():
"""优先读环境变量；否则在脚本所在目录及其父目录找 .env。"""
key = os.environ.get("APIYI_API_KEY")
if key:
return key
here = os.path.dirname(os.path.abspath(__file__))
for d in (here, os.path.dirname(here)):
env_path = os.path.join(d, ".env")
if os.path.exists(env_path):
with open(env_path, encoding="utf-8") as f:
for line in f:
line = line.strip()
if line.startswith("APIYI_API_KEY") and "=" in line:
return line.split("=", 1)[1].strip().strip('"').strip("'")
return None
def project_root():
"""从脚本位置向上找包含 .git 或 .claude 的目录，作为项目根目录；找不到则用当前工作目录。"""
d = os.path.dirname(os.path.abspath(__file__))
while True:
if os.path.isdir(os.path.join(d, ".git")) or os.path.isdir(os.path.join(d, ".claude")):
return d
parent = os.path.dirname(d)
if parent == d:
return os.getcwd()
d = parent
def to_b64(path):
with open(path, "rb") as f:
return base64.b64encode(f.read()).decode()
def mime_of(path):
return "image/png" if path.lower().endswith(".png") else "image/jpeg"
def generate(api_key, endpoint, prompt, images, aspect, size):
"""发一次请求，返回图片字节；失败抛 RuntimeError。"""
parts = [{"text": prompt}]
for path in images:
parts.append({"inlineData": {"mimeType": mime_of(path), "data": to_b64(path)}})
payload = json.dumps({
"contents": [{"parts": parts}],
"generationConfig": {
"responseModalities": ["IMAGE"],
"imageConfig": {"aspectRatio": aspect, "imageSize": size},
},
}).encode()
req = urllib.request.Request(
endpoint, data=payload, method="POST",
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
)
try:
with urllib.request.urlopen(req, timeout=360) as r:
resp = json.loads(r.read())
except urllib.error.HTTPError as e:
raise RuntimeError(f"请求失败 HTTP {e.code}：{e.read().decode(errors='replace')}")
candidates = resp.get("candidates")
if not candidates:
raise RuntimeError(f"未返回候选内容（可能被内容安全策略拒绝）：{resp}")
cand = candidates[0]
# 内容审核拦截：finishReason 非 STOP，或只返回了文字说明
if cand.get("finishReason") not in (None, "STOP"):
text = next((p.get("text") for p in cand["content"]["parts"] if p.get("text")), "")
raise RuntimeError(f"请求被拒绝（finishReason={cand.get('finishReason')}）：{text}")
image_part = next((p for p in cand["content"]["parts"] if p.get("inlineData")), None)
if not image_part:
text = next((p.get("text") for p in cand["content"]["parts"] if p.get("text")), "")
raise RuntimeError(f"未返回图片，模型回复：{text}")
return base64.b64decode(image_part["inlineData"]["data"])
def resolve_paths(out, count):
"""决定输出路径列表。
- 若 out 带目录成分（相对/绝对），按用户给定的路径处理（相对则相对当前工作目录）。
- 若 out 是纯文件名，统一存到 <项目根>/nano-banana-output/ 下，确保好找。
count>1 时给文件名加 -1 / -2 … 后缀。
"""
if os.path.dirname(out):
base_path = os.path.abspath(out)
else:
out_dir = os.path.join(project_root(), "nano-banana-output")
os.makedirs(out_dir, exist_ok=True)
base_path = os.path.join(out_dir, out)
if count == 1:
return [base_path]
base, ext = os.path.splitext(base_path)
return [f"{base}-{i}{ext}" for i in range(1, count + 1)]
def main():
api_key = load_api_key()
if not api_key:
sys.exit("未找到 API Key：请在技能目录的 .env 写一行 APIYI_API_KEY=sk-xxx")
model = os.environ.get("APIYI_IMAGE_MODEL", "gemini-3.1-flash-image-preview")
endpoint = f"https://api.apiyi.com/v1beta/models/{model}:generateContent"
parser = argparse.ArgumentParser(description="Nano Banana 2 出图")
parser.add_argument("prompt", help="提示词 / 编辑指令")
parser.add_argument("-i", "--image", action="append", default=[],
help="输入图片路径（可重复，传入即为编辑模式）")
parser.add_argument("-o", "--out", default="output.png", help="输出文件名")
parser.add_argument("-n", "--count", type=int, default=1,
help=f"一次出几张，默认 1，最多 {MAX_COUNT}（并发生成）")
parser.add_argument("--aspect", default="1:1", help="宽高比，14 选 1，如 16:9 / 1:4 / 8:1")
parser.add_argument("--size", default="2K", help="分辨率 512 / 1K / 2K / 4K")
args = parser.parse_args()
count = args.count
if count < 1:
count = 1
if count > MAX_COUNT:
print(f"提示：一次最多 {MAX_COUNT} 张，已将 {args.count} 限制为 {MAX_COUNT}。", file=sys.stderr)
count = MAX_COUNT
paths = resolve_paths(args.out, count)
def task(path):
data = generate(api_key, endpoint, args.prompt, args.image, args.aspect, args.size)
with open(path, "wb") as f:
f.write(data)
return os.path.abspath(path)
failures = 0
with ThreadPoolExecutor(max_workers=count) as pool:
for path, result in zip(paths, pool.map(lambda p: _safe(task, p), paths)):
ok, value = result
if ok:
print(f"图片已保存至 {value}")
else:
failures += 1
print(f"第 {os.path.basename(path)} 张生成失败：{value}", file=sys.stderr)
if failures == count:
sys.exit("全部生成失败。")
def _safe(fn, arg):
try:
return True, fn(arg)
except Exception as e: # noqa: BLE001 — 单张失败不影响其它并发任务
return False, str(e)
if __name__ == "__main__":
main()
