140 lines
5.6 KiB
Python
140 lines
5.6 KiB
Python
|
"""
|
|||
|
|
图片生成服务 - 使用 Gemini API 生图，存储到本地/云托管
|
|||
|
|
"""
|
|||
|
|
import httpx
|
|||
|
|
import base64
|
|||
|
|
import time
|
|||
|
|
import hashlib
|
|||
|
|
import os
|
|||
|
|
from typing import Optional
|
|||
|
|
|
|||
|
|
from app.config import get_settings
|
|||
|
|
|
|||
|
|
# 图片尺寸规范（基于前端展示尺寸 × 3倍清晰度）
|
|||
|
|
IMAGE_SIZES = {
|
|||
|
|
"cover": {"width": 240, "height": 330, "desc": "竖版封面图，3:4比例"},
|
|||
|
|
"avatar": {"width": 150, "height": 150, "desc": "正方形头像，1:1比例"},
|
|||
|
|
"background": {"width": 1120, "height": 840, "desc": "横版背景图，4:3比例"},
|
|||
|
|
"character": {"width": 512, "height": 768, "desc": "竖版角色立绘，2:3比例，透明背景"}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ImageGenService:
|
|||
|
|
def __init__(self):
|
|||
|
|
settings = get_settings()
|
|||
|
|
self.api_key = settings.gemini_api_key
|
|||
|
|
self.base_url = "https://work.poloapi.com/v1beta"
|
|||
|
|
# 计算绝对路径
|
|||
|
|
base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..', settings.upload_path))
|
|||
|
|
self.upload_dir = os.path.join(base_dir, "images")
|
|||
|
|
os.makedirs(self.upload_dir, exist_ok=True)
|
|||
|
|
|
|||
|
|
def get_size_prompt(self, image_type: str) -> str:
|
|||
|
|
"""获取尺寸描述提示词"""
|
|||
|
|
size = IMAGE_SIZES.get(image_type, IMAGE_SIZES["background"])
|
|||
|
|
return f"Image size: {size['width']}x{size['height']} pixels, {size['desc']}."
|
|||
|
|
|
|||
|
|
async def generate_image(self, prompt: str, image_type: str = "background", style: str = "anime") -> Optional[dict]:
|
|||
|
|
"""
|
|||
|
|
调用 Gemini 生成图片
|
|||
|
|
prompt: 图片描述
|
|||
|
|
image_type: 图片类型（cover/avatar/background/character）
|
|||
|
|
style: 风格（anime/realistic/illustration）
|
|||
|
|
"""
|
|||
|
|
style_prefix = {
|
|||
|
|
"anime": "anime style, high quality illustration, vibrant colors, ",
|
|||
|
|
"realistic": "photorealistic, high detail, cinematic lighting, ",
|
|||
|
|
"illustration": "digital art illustration, beautiful artwork, "
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 组合完整提示词：风格 + 尺寸 + 内容
|
|||
|
|
size_prompt = self.get_size_prompt(image_type)
|
|||
|
|
full_prompt = f"{style_prefix.get(style, '')}{size_prompt} {prompt}"
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
async with httpx.AsyncClient(timeout=60.0) as client:
|
|||
|
|
response = await client.post(
|
|||
|
|
f"{self.base_url}/models/gemini-3-pro-image-preview:generateContent",
|
|||
|
|
headers={
|
|||
|
|
"Content-Type": "application/json",
|
|||
|
|
"x-goog-api-key": self.api_key
|
|||
|
|
},
|
|||
|
|
json={
|
|||
|
|
"contents": [{
|
|||
|
|
"parts": [{
|
|||
|
|
"text": f"Generate an image: {full_prompt}"
|
|||
|
|
}]
|
|||
|
|
}],
|
|||
|
|
"generationConfig": {
|
|||
|
|
"responseModalities": ["TEXT", "IMAGE"]
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if response.status_code == 200:
|
|||
|
|
data = response.json()
|
|||
|
|
candidates = data.get("candidates", [])
|
|||
|
|
if candidates:
|
|||
|
|
parts = candidates[0].get("content", {}).get("parts", [])
|
|||
|
|
for part in parts:
|
|||
|
|
if "inlineData" in part:
|
|||
|
|
return {
|
|||
|
|
"success": True,
|
|||
|
|
"image_data": part["inlineData"]["data"],
|
|||
|
|
"mime_type": part["inlineData"].get("mimeType", "image/png")
|
|||
|
|
}
|
|||
|
|
return {"success": False, "error": "No image in response"}
|
|||
|
|
else:
|
|||
|
|
error_text = response.text[:200]
|
|||
|
|
return {"success": False, "error": f"API error: {response.status_code} - {error_text}"}
|
|||
|
|
except Exception as e:
|
|||
|
|
return {"success": False, "error": str(e)}
|
|||
|
|
|
|||
|
|
async def save_image(self, image_data: str, filename: str) -> Optional[str]:
|
|||
|
|
"""保存图片到本地，返回访问URL"""
|
|||
|
|
try:
|
|||
|
|
image_bytes = base64.b64decode(image_data)
|
|||
|
|
file_path = os.path.join(self.upload_dir, filename)
|
|||
|
|
|
|||
|
|
with open(file_path, "wb") as f:
|
|||
|
|
f.write(image_bytes)
|
|||
|
|
|
|||
|
|
# 返回可访问的URL路径
|
|||
|
|
return f"/uploads/images/{filename}"
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"Save error: {e}")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
async def generate_and_save(self, prompt: str, image_type: str = "background", style: str = "anime") -> dict:
|
|||
|
|
"""生成图片并保存"""
|
|||
|
|
result = await self.generate_image(prompt, image_type, style)
|
|||
|
|
|
|||
|
|
if not result or not result.get("success"):
|
|||
|
|
return {
|
|||
|
|
"success": False,
|
|||
|
|
"error": result.get("error", "生成失败") if result else "生成失败"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 生成文件名
|
|||
|
|
timestamp = int(time.time() * 1000)
|
|||
|
|
hash_str = hashlib.md5(prompt.encode()).hexdigest()[:8]
|
|||
|
|
ext = "png" if "png" in result.get("mime_type", "") else "jpg"
|
|||
|
|
filename = f"{image_type}_{timestamp}_{hash_str}.{ext}"
|
|||
|
|
|
|||
|
|
url = await self.save_image(result["image_data"], filename)
|
|||
|
|
|
|||
|
|
if url:
|
|||
|
|
return {"success": True, "url": url, "filename": filename}
|
|||
|
|
else:
|
|||
|
|
return {"success": False, "error": "保存失败"}
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 延迟初始化单例
|
|||
|
|
_service_instance = None
|
|||
|
|
|
|||
|
|
def get_image_gen_service():
|
|||
|
|
global _service_instance
|
|||
|
|
if _service_instance is None:
|
|||
|
|
_service_instance = ImageGenService()
|
|||
|
|
return _service_instance
|