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
|