starmen1265
7fbef405ba
image_gen.py: --fuse支持1风格+N张内容图(-c可重复); 新增-pf/--prompt-file从文件读prompt; 修复images.edit参数格式bug(dict->tuple); save_result改用短英文文件名(修复长中文文件名导致图像分析工具URL解析失败). README.md: 工具链概述文档.
481 lines
18 KiB
Python
481 lines
18 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
ChildHood AI Art Generation Tool
|
||
通过 API易 平台调用 AI 生图模型
|
||
|
||
Usage:
|
||
python image_gen.py --test # 测试 API 连通性
|
||
python image_gen.py --generate "千禧年中国城郊的小卖部门口" # 文生图（自动加风格前缀）
|
||
python image_gen.py --generate "a red apple" --no-style # 纯文生图（不加风格前缀）
|
||
python image_gen.py --calibrate # 风格校准（5 个测试物体）
|
||
python image_gen.py --fuse --style style.png --content photo.jpg # 双参考融合
|
||
python image_gen.py --fuse --style style.png --content photo.jpg --prompt "自定义 prompt"
|
||
|
||
Requirements:
|
||
pip install openai python-dotenv
|
||
"""
|
||
|
||
import argparse
|
||
import base64
|
||
import io
|
||
import json
|
||
import os
|
||
import sys
|
||
import urllib.request
|
||
|
||
# Windows 控制台 UTF-8 支持
|
||
if sys.platform == "win32":
|
||
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")
|
||
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding="utf-8", errors="replace")
|
||
import os
|
||
import sys
|
||
import urllib.request
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 路径配置
|
||
# ---------------------------------------------------------------------------
|
||
SCRIPT_DIR = Path(__file__).resolve().parent # ai-art-research/tools/
|
||
RESEARCH_DIR = SCRIPT_DIR.parent # ai-art-research/
|
||
DOTENV_PATH = RESEARCH_DIR / ".env"
|
||
OUTPUT_BASE = RESEARCH_DIR / "生成结果"
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 加载 .env
|
||
# ---------------------------------------------------------------------------
|
||
if DOTENV_PATH.exists():
|
||
from dotenv import load_dotenv
|
||
load_dotenv(DOTENV_PATH)
|
||
|
||
API_KEY = os.getenv("APIYI_API_KEY", "")
|
||
BASE_URL = os.getenv("APIYI_BASE_URL", "https://api.apiyi.com/v1")
|
||
DEFAULT_MODEL = os.getenv("APIYI_MODEL", "gpt-image-2")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 风格 Prompt 种子（来自工作流方案 v0.1）
|
||
# ---------------------------------------------------------------------------
|
||
STYLE_SEED_EN = (
|
||
"3D low-poly game art, flat cel-shaded with solid color blocks, "
|
||
"bright warm-cool balanced saturated palette, soft diffused lighting "
|
||
"with subtle shadows, simplified geometric forms with clean edges, "
|
||
"low-angle third-person perspective, nostalgic warm summer atmosphere"
|
||
)
|
||
|
||
# 风格校准测试物体（中文名, 英文 prompt）
|
||
CALIBRATION_OBJECTS = [
|
||
("乌鸦", "a crow standing on the ground"),
|
||
("茶杯", "a teacup on a wooden table"),
|
||
("老人", "an old man standing in a village street"),
|
||
("山湖", "a mountain lake with trees"),
|
||
("苹果", "a red apple on the ground"),
|
||
]
|
||
|
||
# 双参考融合默认 prompt
|
||
DEFAULT_FUSION_PROMPT = """\
|
||
图1 是我的风格参考（一张游戏截图）。图2 是我的内容参考（一张现实照片）。
|
||
请生成一张新图片，用图1的美术风格来呈现图2中的场景内容。
|
||
|
||
风格要求：
|
||
- 高度简化的低多边形几何形体，干净利落的边缘
|
||
- 平面色块着色，不用渐变或写实纹理
|
||
- 鲜艳明亮的暖冷平衡色调（暖黄、天蓝、翠绿为主）
|
||
- 柔和散射光，阴影淡而柔和
|
||
- 微低角度第三人称视角
|
||
- 温暖怀旧的夏日午后氛围
|
||
|
||
内容要求：
|
||
- 保留图2中的建筑类型、布局、空间关系
|
||
- 但把所有建筑元素替换为千禧年中国城郊的风格
|
||
- 色调保持温暖怀旧"""
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 工具函数
|
||
# ---------------------------------------------------------------------------
|
||
def get_client() -> "OpenAI":
|
||
"""创建 OpenAI 客户端"""
|
||
if not API_KEY:
|
||
print("❌ 未找到 API Key")
|
||
print(f" 请在 {DOTENV_PATH} 中设置: APIYI_API_KEY=sk-xxx")
|
||
sys.exit(1)
|
||
from openai import OpenAI
|
||
return OpenAI(api_key=API_KEY, base_url=BASE_URL)
|
||
|
||
|
||
def _extract_image_bytes(img_obj) -> bytes:
|
||
"""从 API 响应中提取图片字节数据"""
|
||
# 优先 base64
|
||
b64 = getattr(img_obj, "b64_json", None)
|
||
if b64:
|
||
return base64.b64decode(b64)
|
||
# 其次 URL
|
||
url = getattr(img_obj, "url", None)
|
||
if url:
|
||
print(" ⬇️ 正在下载图片...")
|
||
return urllib.request.urlopen(url).read()
|
||
raise ValueError("响应中未找到图片数据（无 b64_json 也无 url）")
|
||
|
||
|
||
def _sanitize_filename(s: str, max_len: int = 40) -> str:
|
||
"""清理文件名中的非法字符"""
|
||
result = "".join(c if c.isalnum() or c in "._-" else "_" for c in s)
|
||
return result[:max_len].rstrip("_")
|
||
|
||
|
||
# 模式 → 简短英文文件名 tag
|
||
# 避免把整段中文 prompt 塞进文件名（CDN URL 过长/编码异常会让图像分析工具解析失败）
|
||
_MODE_SLUG = {
|
||
"text_to_image": "gen",
|
||
"style_calibration": "calib",
|
||
"multi_reference_fusion": "fusion",
|
||
"dual_reference_fusion": "fusion",
|
||
"connection_test": "test",
|
||
}
|
||
|
||
|
||
def save_result(image_bytes: bytes, metadata: dict, subdir: str = "") -> Path:
|
||
"""保存图片和元数据，返回图片路径。
|
||
文件名用「时间戳_模型_模式」短英文命名（如 260617-183020_gpt-image-2_fusion.png）；
|
||
完整 prompt 与参数存在 .meta.json 里，不丢信息。
|
||
"""
|
||
ts = datetime.now().strftime("%y%m%d-%H%M%S")
|
||
model_slug = metadata.get("model", "unknown").replace("/", "-").replace(".", "-")
|
||
mode_slug = _MODE_SLUG.get(metadata.get("mode", ""), "img")
|
||
|
||
out_dir = OUTPUT_BASE / subdir if subdir else OUTPUT_BASE
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
filename = f"{ts}_{model_slug}_{mode_slug}.png"
|
||
|
||
img_path = out_dir / filename
|
||
img_path.write_bytes(image_bytes)
|
||
|
||
meta_path = img_path.with_suffix(".meta.json")
|
||
meta_path.write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
|
||
print(f"✅ 已保存: {img_path.name}")
|
||
return img_path
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 功能：测试连通性
|
||
# ---------------------------------------------------------------------------
|
||
def cmd_test(client):
|
||
"""测试 API 易连通性，生成一张简单测试图"""
|
||
print("🔗 测试 API易 连通性...")
|
||
print(f" Base URL: {BASE_URL}")
|
||
print(f" Model: {DEFAULT_MODEL}")
|
||
print(f" API Key: {API_KEY[:8]}...{API_KEY[-4:]}")
|
||
print()
|
||
|
||
prompt = f"{STYLE_SEED_EN}, a simple red cube on a white background, test image"
|
||
print(f" Prompt: {prompt[:80]}...")
|
||
|
||
response = client.images.generate(
|
||
model=DEFAULT_MODEL,
|
||
prompt=prompt,
|
||
size="1024x1024",
|
||
quality="low",
|
||
n=1,
|
||
)
|
||
|
||
image_bytes = _extract_image_bytes(response.data[0])
|
||
save_result(image_bytes, {
|
||
"model": DEFAULT_MODEL,
|
||
"prompt": prompt,
|
||
"mode": "connection_test",
|
||
"size": "1024x1024",
|
||
"quality": "low",
|
||
"timestamp": datetime.now().isoformat(),
|
||
}, subdir="测试")
|
||
|
||
print()
|
||
print("✅ API 连通性测试通过！")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 功能：文生图
|
||
# ---------------------------------------------------------------------------
|
||
def cmd_generate(client, prompt: str, use_style: bool, size: str, quality: str,
|
||
model: str, subdir: str):
|
||
"""文生图"""
|
||
model = model or DEFAULT_MODEL
|
||
full_prompt = f"{STYLE_SEED_EN}, {prompt}" if use_style else prompt
|
||
|
||
print(f"🎨 文生图")
|
||
print(f" 模型: {model}")
|
||
print(f" 尺寸: {size} 质量: {quality}")
|
||
print(f" 风格前缀: {'是' if use_style else '否'}")
|
||
print(f" Prompt: {full_prompt[:100]}{'...' if len(full_prompt) > 100 else ''}")
|
||
print()
|
||
|
||
response = client.images.generate(
|
||
model=model,
|
||
prompt=full_prompt,
|
||
size=size,
|
||
quality=quality,
|
||
n=1,
|
||
)
|
||
|
||
image_bytes = _extract_image_bytes(response.data[0])
|
||
save_result(image_bytes, {
|
||
"model": model,
|
||
"prompt": full_prompt,
|
||
"mode": "text_to_image",
|
||
"style_prefix_used": use_style,
|
||
"size": size,
|
||
"quality": quality,
|
||
"timestamp": datetime.now().isoformat(),
|
||
}, subdir=subdir)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 功能：风格校准
|
||
# ---------------------------------------------------------------------------
|
||
def cmd_calibrate(client, model: str):
|
||
"""风格校准 — 5 个测试物体"""
|
||
model = model or DEFAULT_MODEL
|
||
subdir = "风格校准"
|
||
|
||
print("🎯 开始风格校准测试")
|
||
print(f" 模型: {model}")
|
||
print(f" 风格种子: {STYLE_SEED_EN[:60]}...")
|
||
print(f" 测试物体: {', '.join(zh for zh, _ in CALIBRATION_OBJECTS)}")
|
||
print()
|
||
|
||
success = 0
|
||
for zh_name, en_prompt in CALIBRATION_OBJECTS:
|
||
full_prompt = f"{STYLE_SEED_EN}, {en_prompt}"
|
||
print(f" 🎨 生成: {zh_name} ...")
|
||
|
||
try:
|
||
response = client.images.generate(
|
||
model=model,
|
||
prompt=full_prompt,
|
||
size="1024x1024",
|
||
quality="medium",
|
||
n=1,
|
||
response_format="b64_json",
|
||
)
|
||
image_bytes = _extract_image_bytes(response.data[0])
|
||
save_result(image_bytes, {
|
||
"model": model,
|
||
"prompt": full_prompt,
|
||
"mode": "style_calibration",
|
||
"object_zh": zh_name,
|
||
"object_en": en_prompt,
|
||
"style_seed": STYLE_SEED_EN,
|
||
"timestamp": datetime.now().isoformat(),
|
||
}, subdir=subdir)
|
||
success += 1
|
||
except Exception as e:
|
||
print(f" ❌ 失败: {e}")
|
||
print()
|
||
|
||
print(f"🎉 风格校准完成！成功 {success}/{len(CALIBRATION_OBJECTS)} 张")
|
||
print(f" 输出目录: {OUTPUT_BASE / subdir}")
|
||
print()
|
||
print("📋 验证方法：打开上述目录，检查 5 张图是否视觉风格统一。")
|
||
print(" 如果风格不一致，需要调整 STYLE_SEED_EN 后重新运行。")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 功能：双参考融合
|
||
# ---------------------------------------------------------------------------
|
||
def cmd_fuse(client, style_path: str, content_paths, prompt: str,
|
||
model: str, subdir: str):
|
||
"""参考图融合（1 张风格参考 + 1~N 张内容参考）"""
|
||
model = model or DEFAULT_MODEL
|
||
fusion_prompt = prompt or DEFAULT_FUSION_PROMPT
|
||
subdir = subdir or "双参考融合"
|
||
|
||
style_file = Path(style_path).resolve()
|
||
if not style_file.exists():
|
||
print(f"❌ 风格参考图不存在: {style_file}")
|
||
sys.exit(1)
|
||
|
||
# content_paths 兼容单个 str/Path 或 list
|
||
if isinstance(content_paths, (str, Path)):
|
||
content_paths = [content_paths]
|
||
content_files = []
|
||
for cp in content_paths:
|
||
cf = Path(cp).resolve()
|
||
if not cf.exists():
|
||
print(f"❌ 内容参考图不存在: {cf}")
|
||
sys.exit(1)
|
||
content_files.append(cf)
|
||
|
||
label = "双参考融合" if len(content_files) == 1 else f"多参考融合（1 风格 + {len(content_files)} 内容）"
|
||
print(f"🔗 {label}")
|
||
print(f" 模型: {model}")
|
||
print(f" 风格参考: {style_file.name}")
|
||
for cf in content_files:
|
||
print(f" 内容参考: {cf.name}")
|
||
print(f" Prompt: {fusion_prompt[:80]}...")
|
||
print()
|
||
|
||
# 确定图片 MIME 类型
|
||
ext_map = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg",
|
||
".webp": "image/webp", ".gif": "image/gif"}
|
||
|
||
# openai 库文件上传格式：tuple = (filename, bytes, content_type)
|
||
style_tuple = (style_file.name, style_file.read_bytes(),
|
||
ext_map.get(style_file.suffix.lower(), "image/png"))
|
||
content_tuples = [
|
||
(cf.name, cf.read_bytes(), ext_map.get(cf.suffix.lower(), "image/png"))
|
||
for cf in content_files
|
||
]
|
||
image_list = [style_tuple] + content_tuples
|
||
|
||
image_bytes = None
|
||
errors = []
|
||
|
||
# 方式 1: images.edit 多图（风格参考 + 全部内容参考）
|
||
print(f" 尝试方式 1: images.edit ({len(image_list)} 图) 融合 ...")
|
||
try:
|
||
response = client.images.edit(
|
||
model=model,
|
||
image=image_list,
|
||
prompt=fusion_prompt,
|
||
size="1024x1024",
|
||
)
|
||
image_bytes = _extract_image_bytes(response.data[0])
|
||
print(" ✅ 方式 1 成功")
|
||
except Exception as e:
|
||
errors.append(f"images.edit 多图: {e}")
|
||
print(f" ⚠️ 方式 1 失败: {e}")
|
||
|
||
# 方式 2: images.edit 单图（仅第一张内容参考，风格靠文字）—— 降级
|
||
if image_bytes is None:
|
||
print(" 尝试方式 2: images.edit 单图（风格写入 prompt）...")
|
||
single_prompt = (
|
||
f"{STYLE_SEED_EN}\n\n"
|
||
"Re-render the uploaded photo into this 3D low-poly cel-shaded game art style. "
|
||
"Keep the photo's scene layout, building structure and objects, but re-render "
|
||
"everything in the style described above. Warm nostalgic late-afternoon lighting, "
|
||
"clean silhouette, game concept art."
|
||
)
|
||
try:
|
||
response = client.images.edit(
|
||
model=model,
|
||
image=content_tuples[0],
|
||
prompt=single_prompt,
|
||
size="1024x1024",
|
||
)
|
||
image_bytes = _extract_image_bytes(response.data[0])
|
||
print(" ✅ 方式 2 成功（单图 + 文字风格，效果弱于多图）")
|
||
except Exception as e:
|
||
errors.append(f"images.edit 单图: {e}")
|
||
print(f" ⚠️ 方式 2 失败: {e}")
|
||
|
||
if image_bytes is None:
|
||
print()
|
||
print("❌ 双参考融合 API 调用失败（双图 edit 与单图 edit 均失败）。")
|
||
print()
|
||
print("💡 可能原因与替代方案：")
|
||
print(" 1. API易 的该模型可能不支持 images.edit → 改用网页端手动上传两张图")
|
||
print(" 2. 换支持参考图的模型重试: -m gemini-3-1-flash-image-preview")
|
||
print(" 3. 改走纯文生图: python tools/image_gen.py -g \"<建筑内容描述>\"")
|
||
print()
|
||
print(" 所有错误:")
|
||
for err in errors:
|
||
print(f" - {err}")
|
||
sys.exit(1)
|
||
|
||
save_result(image_bytes, {
|
||
"model": model,
|
||
"prompt": fusion_prompt,
|
||
"mode": "multi_reference_fusion",
|
||
"style_reference": str(style_file),
|
||
"content_references": [str(cf) for cf in content_files],
|
||
"size": "1024x1024",
|
||
"timestamp": datetime.now().isoformat(),
|
||
}, subdir=subdir)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CLI 入口
|
||
# ---------------------------------------------------------------------------
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description="ChildHood AI Art Generation Tool — 通过 API易 生图",
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
epilog="""\
|
||
示例:
|
||
python image_gen.py --test # 测试连通性
|
||
python image_gen.py -g "千禧年中国城郊的小卖部门口" # 文生图
|
||
python image_gen.py -g "a red apple" --no-style # 不加风格前缀
|
||
python image_gen.py --calibrate # 风格校准
|
||
python image_gen.py --fuse --style ref.png --content photo.jpg # 双参考融合
|
||
python image_gen.py --fuse -s style.png -c photo.jpg -p "自定义" # 自定义 prompt
|
||
|
||
环境变量 (ai-art-research/.env):
|
||
APIYI_API_KEY API易 令牌 (必填)
|
||
APIYI_BASE_URL API 地址 (默认 https://api.apiyi.com/v1)
|
||
APIYI_MODEL 默认模型 (默认 gpt-image-2)
|
||
""",
|
||
)
|
||
|
||
# 互斥的操作模式
|
||
mode = parser.add_mutually_exclusive_group()
|
||
mode.add_argument("--test", action="store_true", help="测试 API 连通性")
|
||
mode.add_argument("-g", "--generate", type=str, metavar="PROMPT", help="文生图")
|
||
mode.add_argument("--calibrate", action="store_true", help="风格校准（5 个测试物体）")
|
||
mode.add_argument("--fuse", action="store_true", help="双参考图融合模式")
|
||
|
||
# 参数
|
||
parser.add_argument("-s", "--style", type=str, help="风格参考图路径（配合 --fuse）")
|
||
parser.add_argument("-c", "--content", action="append", type=str,
|
||
help="内容参考图路径（配合 --fuse，可重复传多张：首张为整楼，后续为细节特写如晾衣架）")
|
||
parser.add_argument("-p", "--prompt", type=str, help="自定义 prompt（短）")
|
||
parser.add_argument("-pf", "--prompt-file", type=str,
|
||
help="从 .txt 文件读取 prompt（配合 --fuse，适合长 prompt，避免命令行转义噪音）")
|
||
parser.add_argument("-m", "--model", type=str, help=f"模型 (默认: {DEFAULT_MODEL})")
|
||
parser.add_argument("--no-style", action="store_true", help="文生图不加风格前缀")
|
||
parser.add_argument("--size", type=str, default="1024x1024", help="图片尺寸")
|
||
parser.add_argument("--quality", type=str, default="medium",
|
||
choices=["low", "medium", "high"], help="质量")
|
||
parser.add_argument("--subdir", type=str, default="", help="输出子目录名")
|
||
|
||
args = parser.parse_args()
|
||
|
||
# 无参数时显示帮助
|
||
if len(sys.argv) == 1:
|
||
parser.print_help()
|
||
return
|
||
|
||
client = get_client()
|
||
|
||
if args.test:
|
||
cmd_test(client)
|
||
elif args.generate:
|
||
cmd_generate(
|
||
client, args.generate,
|
||
use_style=not args.no_style,
|
||
size=args.size, quality=args.quality,
|
||
model=args.model, subdir=args.subdir or "文生图",
|
||
)
|
||
elif args.calibrate:
|
||
cmd_calibrate(client, model=args.model)
|
||
elif args.fuse:
|
||
if not args.style or not args.content:
|
||
parser.error("双参考融合需要 --style (-s) 和 --content (-c) 参数")
|
||
if args.prompt_file:
|
||
pf = Path(args.prompt_file)
|
||
if not pf.exists():
|
||
parser.error(f"prompt 文件不存在: {pf}")
|
||
fuse_prompt = pf.read_text(encoding="utf-8")
|
||
print(f" 📄 从文件读取 prompt: {pf.name}")
|
||
else:
|
||
fuse_prompt = args.prompt
|
||
cmd_fuse(
|
||
client, args.style, args.content,
|
||
prompt=fuse_prompt, model=args.model, subdir=args.subdir,
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|