You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
114 lines
3.8 KiB
114 lines
3.8 KiB
#!/usr/bin/env python3
|
|
"""通过 FastAPI 提供 GPT Image 和 Grok Image 文生图服务。"""
|
|
|
|
import base64
|
|
import os
|
|
import time
|
|
from pathlib import Path
|
|
from uuid import uuid4
|
|
|
|
from dotenv import load_dotenv
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.staticfiles import StaticFiles
|
|
from openai import APIConnectionError, APIStatusError, APITimeoutError, OpenAI
|
|
from pydantic import BaseModel, Field
|
|
|
|
load_dotenv()
|
|
|
|
API_BASE_URL = os.getenv("IMAGE_API_BASE_URL", "https://api.slomerex.xyz/v1")
|
|
OUTPUT_DIR = Path("output")
|
|
MAX_RETRIES = 3
|
|
RETRYABLE_STATUS_CODES = {408, 429, 500, 502, 503, 504}
|
|
PROVIDERS = {
|
|
"grok": {
|
|
"api_key_env": "GROK_API_KEY",
|
|
"model": "grok-imagine-image-lite",
|
|
},
|
|
"gpt": {
|
|
"api_key_env": "GPT_API_KEY",
|
|
"model": "gpt-image-2-1K",
|
|
},
|
|
}
|
|
|
|
app = FastAPI(title="AI Image T2I API", version="1.0.0")
|
|
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
|
app.mount("/output", StaticFiles(directory=OUTPUT_DIR), name="output")
|
|
|
|
|
|
class ImageGenerationRequest(BaseModel):
|
|
prompt: str = Field(min_length=1, description="图片提示词")
|
|
provider: str = Field(default="grok", pattern="^(grok|gpt)$", description="模型提供方")
|
|
model: str | None = Field(default=None, description="覆盖提供方默认模型")
|
|
size: str = Field(default="1024x1024", description="图片尺寸")
|
|
n: int = Field(default=1, ge=1, le=10, description="生成数量")
|
|
|
|
|
|
class GeneratedImage(BaseModel):
|
|
filename: str
|
|
url: str
|
|
bytes: int
|
|
|
|
|
|
class ImageGenerationResponse(BaseModel):
|
|
provider: str
|
|
model: str
|
|
images: list[GeneratedImage]
|
|
|
|
|
|
def generate_with_retry(client: OpenAI, *, model: str, prompt: str, size: str, n: int):
|
|
for attempt in range(1, MAX_RETRIES + 1):
|
|
try:
|
|
return client.images.generate(model=model, prompt=prompt, size=size, n=n)
|
|
except APIStatusError as e:
|
|
if e.status_code not in RETRYABLE_STATUS_CODES or attempt == MAX_RETRIES:
|
|
raise HTTPException(status_code=502, detail=f"上游图片服务错误: HTTP {e.status_code}") from e
|
|
time.sleep(2 ** (attempt - 1))
|
|
except (APIConnectionError, APITimeoutError) as e:
|
|
if attempt == MAX_RETRIES:
|
|
raise HTTPException(status_code=504, detail="上游图片服务连接或请求超时") from e
|
|
time.sleep(2 ** (attempt - 1))
|
|
|
|
|
|
@app.get("/health")
|
|
def health_check():
|
|
return {"status": "ok"}
|
|
|
|
|
|
@app.post("/images/generations", response_model=ImageGenerationResponse)
|
|
def generate_images(request: ImageGenerationRequest):
|
|
config = PROVIDERS[request.provider]
|
|
api_key = os.getenv(config["api_key_env"])
|
|
if not api_key:
|
|
raise HTTPException(status_code=500, detail=f"未配置环境变量 {config['api_key_env']}")
|
|
|
|
model = request.model or config["model"]
|
|
client = OpenAI(base_url=API_BASE_URL, api_key=api_key, timeout=120.0, max_retries=0)
|
|
response = generate_with_retry(
|
|
client,
|
|
model=model,
|
|
prompt=request.prompt,
|
|
size=request.size,
|
|
n=request.n,
|
|
)
|
|
|
|
images = []
|
|
for item in response.data:
|
|
if not item.b64_json:
|
|
raise HTTPException(status_code=502, detail="上游图片服务未返回 Base64 图片数据")
|
|
|
|
try:
|
|
raw = base64.b64decode(item.b64_json)
|
|
except ValueError as e:
|
|
raise HTTPException(status_code=502, detail="上游图片数据格式无效") from e
|
|
|
|
filename = f"{model}-{uuid4().hex}.png"
|
|
(OUTPUT_DIR / filename).write_bytes(raw)
|
|
images.append(GeneratedImage(filename=filename, url=f"/output/{filename}", bytes=len(raw)))
|
|
|
|
return ImageGenerationResponse(provider=request.provider, model=model, images=images)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
|
|
uvicorn.run(app, host="0.0.0.0", port=8000)
|
|
|