main
parent
0d63239270
commit
fdbc22441b
@ -1,3 +1,7 @@ |
||||
GROK_API_KEY=your_grok_api_key |
||||
GPT_API_KEY=your_gpt_api_key |
||||
IMAGE_API_BASE_URL=https://api.slomerex.xyz/v1 |
||||
|
||||
# Agnes Image 2.1 Flash |
||||
AGNES_API_KEY=your_agnes_api_key |
||||
AGNES_API_BASE_URL=https://apihub.agnes-ai.com/v1 |
||||
|
||||
File diff suppressed because one or more lines are too long
@ -0,0 +1,103 @@ |
||||
#!/usr/bin/env python3 |
||||
"""FastAPI 入口：AI 图片生成服务。""" |
||||
|
||||
from dotenv import load_dotenv |
||||
|
||||
load_dotenv() |
||||
|
||||
from fastapi import FastAPI |
||||
from fastapi.staticfiles import StaticFiles |
||||
from pydantic import BaseModel, Field, model_validator |
||||
|
||||
from utils.mm_api_t2i import OUTPUT_DIR, PROVIDERS, generate_images |
||||
from utils.agnes_image import VALID_RATIOS, generate_images as agnes_generate_images |
||||
|
||||
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="生成数量") |
||||
profile_name: str = Field(pattern=r"^[\w-]+$", description="调用方标识（必填），图片会保存到 output/<profile_name>/ 下") |
||||
|
||||
@model_validator(mode="before") |
||||
@classmethod |
||||
def check_profile_name(cls, data): |
||||
if isinstance(data, dict) and not data.get("profile_name"): |
||||
raise ValueError("缺少 profile_name 参数，请在请求中带上你的 agent/profile 名称") |
||||
return data |
||||
|
||||
|
||||
class GeneratedImage(BaseModel): |
||||
filename: str |
||||
url: str |
||||
bytes: int |
||||
|
||||
|
||||
class ImageGenerationResponse(BaseModel): |
||||
provider: str |
||||
model: str |
||||
images: list[GeneratedImage] |
||||
|
||||
|
||||
@app.get("/health") |
||||
def health_check(): |
||||
return {"status": "ok"} |
||||
|
||||
|
||||
@app.post("/images/generations", response_model=ImageGenerationResponse) |
||||
def generate_images_endpoint(request: ImageGenerationRequest): |
||||
result = generate_images( |
||||
provider=request.provider, |
||||
model=request.model, |
||||
prompt=request.prompt, |
||||
size=request.size, |
||||
n=request.n, |
||||
profile_name=request.profile_name, |
||||
) |
||||
return ImageGenerationResponse(**result) |
||||
|
||||
|
||||
class AgnesImageGenerationRequest(BaseModel): |
||||
prompt: str = Field(min_length=1, description="图片提示词") |
||||
size: str = Field(default="1K", description="输出尺寸档位: 1K / 2K / 3K / 4K，也支持精确尺寸如 1024x1024") |
||||
ratio: str | None = Field(default=None, description=f"宽高比，支持 {sorted(VALID_RATIOS)}；与 size 档位配合使用") |
||||
images: list[str] | None = Field(default=None, description="图生图 / 多图合成的输入图像 URL 或 Data URI Base64") |
||||
profile_name: str = Field(pattern=r"^[\w-]+$", description="调用方标识（必填），图片会保存到 output/<profile_name>/ 下") |
||||
|
||||
@model_validator(mode="before") |
||||
@classmethod |
||||
def check_profile_name(cls, data): |
||||
if isinstance(data, dict) and not data.get("profile_name"): |
||||
raise ValueError("缺少 profile_name 参数，请在请求中带上你的 agent/profile 名称") |
||||
return data |
||||
|
||||
|
||||
class AgnesImageGenerationResponse(BaseModel): |
||||
model: str |
||||
size: str |
||||
ratio: str |
||||
images: list[GeneratedImage] |
||||
|
||||
|
||||
@app.post("/agnes/images/generations", response_model=AgnesImageGenerationResponse) |
||||
def agnes_generate_images_endpoint(request: AgnesImageGenerationRequest): |
||||
result = agnes_generate_images( |
||||
prompt=request.prompt, |
||||
size=request.size, |
||||
ratio=request.ratio, |
||||
images=request.images, |
||||
profile_name=request.profile_name, |
||||
) |
||||
return AgnesImageGenerationResponse(**result) |
||||
|
||||
|
||||
if __name__ == "__main__": |
||||
import uvicorn |
||||
|
||||
uvicorn.run(app, host="0.0.0.0", port=8000) |
||||
@ -0,0 +1,140 @@ |
||||
"""Agnes Image 2.1 Flash 文生图 / 图生图 / 多图合成的核心逻辑。""" |
||||
|
||||
import base64 |
||||
import json |
||||
import os |
||||
import time |
||||
import urllib.error |
||||
import urllib.request |
||||
from pathlib import Path |
||||
from uuid import uuid4 |
||||
|
||||
from fastapi import HTTPException |
||||
|
||||
API_BASE_URL = os.getenv("AGNES_API_BASE_URL", "https://apihub.agnes-ai.com/v1") |
||||
OUTPUT_DIR = Path("output") |
||||
MAX_RETRIES = 3 |
||||
RETRYABLE_STATUS_CODES = {408, 429, 500, 502, 503, 504} |
||||
DEFAULT_MODEL = "agnes-image-2.1-flash" |
||||
VALID_RATIOS = {"1:1", "3:4", "4:3", "16:9", "9:16", "2:3", "3:2", "21:9"} |
||||
RATIO_DIMENSIONS = { |
||||
"1:1": {"1K": "1024x1024", "2K": "2048x2048", "3K": "3072x3072", "4K": "4096x4096"}, |
||||
"3:4": {"1K": "864x1152", "2K": "1728x2304", "3K": "2592x3456", "4K": "3456x4608"}, |
||||
"4:3": {"1K": "1152x864", "2K": "2304x1728", "3K": "3456x2592", "4K": "4608x3456"}, |
||||
"16:9": {"1K": "1312x736", "2K": "2624x1472", "3K": "3936x2208", "4K": "5248x2944"}, |
||||
"9:16": {"1K": "736x1312", "2K": "1472x2624", "3K": "2208x3936", "4K": "2944x5248"}, |
||||
"2:3": {"1K": "832x1248", "2K": "1664x2496", "3K": "2496x3744", "4K": "3328x4992"}, |
||||
"3:2": {"1K": "1248x832", "2K": "2496x1664", "3K": "3744x2496", "4K": "4992x3328"}, |
||||
"21:9": {"1K": "1568x672", "2K": "3136x1344", "3K": "4704x2016", "4K": "6272x2688"}, |
||||
} |
||||
|
||||
|
||||
def _resolve_size(size: str, ratio: str | None) -> str: |
||||
"""把 size 档位 (1K/2K/3K/4K) 配合 ratio 解析成实际像素尺寸，方便保存文件。""" |
||||
if ratio and ratio not in VALID_RATIOS: |
||||
raise HTTPException(status_code=400, detail=f"不支持的宽高比: {ratio}") |
||||
if size in {"1K", "2K", "3K", "4K"}: |
||||
if not ratio: |
||||
ratio = "1:1" |
||||
return RATIO_DIMENSIONS[ratio][size] |
||||
return size |
||||
|
||||
|
||||
def _build_payload(*, prompt: str, size: str, ratio: str | None, images: list[str] | None) -> dict: |
||||
"""组装请求体。注意 response_format 必须放在 extra_body 内。""" |
||||
payload: dict = { |
||||
"model": DEFAULT_MODEL, |
||||
"prompt": prompt, |
||||
"size": size, |
||||
} |
||||
extra_body: dict = {} |
||||
if ratio: |
||||
payload["ratio"] = ratio |
||||
if images: |
||||
extra_body["image"] = images |
||||
if extra_body: |
||||
payload["extra_body"] = extra_body |
||||
return payload |
||||
|
||||
|
||||
def _post_with_retry(url: str, headers: dict, payload: dict) -> dict: |
||||
"""调用 Agnes Image 接口，并对可恢复的上游错误进行指数退避重试。""" |
||||
data = json.dumps(payload).encode("utf-8") |
||||
for attempt in range(1, MAX_RETRIES + 1): |
||||
request = urllib.request.Request(url, data=data, headers=headers, method="POST") |
||||
try: |
||||
with urllib.request.urlopen(request, timeout=360) as response: |
||||
return json.loads(response.read().decode("utf-8")) |
||||
except urllib.error.HTTPError as e: |
||||
if e.code not in RETRYABLE_STATUS_CODES or attempt == MAX_RETRIES: |
||||
body = e.read().decode("utf-8", errors="replace") |
||||
raise HTTPException(status_code=502, detail=f"Agnes 图片服务错误: HTTP {e.code} {body}") from e |
||||
except (urllib.error.URLError, TimeoutError) as e: |
||||
if attempt == MAX_RETRIES: |
||||
raise HTTPException(status_code=504, detail="Agnes 图片服务连接或请求超时") from e |
||||
time.sleep(2 ** (attempt - 1)) |
||||
|
||||
|
||||
def _download_image(url: str) -> bytes: |
||||
"""从 URL 或 data URI 下载图片。""" |
||||
if url.startswith("data:image/"): |
||||
try: |
||||
_, encoded = url.split(",", 1) |
||||
return base64.b64decode(encoded, validate=True) |
||||
except (ValueError, TypeError) as e: |
||||
raise HTTPException(status_code=502, detail="图片 data URL 格式无效") from e |
||||
|
||||
request = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"}) |
||||
for attempt in range(1, MAX_RETRIES + 1): |
||||
try: |
||||
with urllib.request.urlopen(request, timeout=120) as response: |
||||
return response.read() |
||||
except urllib.error.HTTPError as e: |
||||
if e.code not in RETRYABLE_STATUS_CODES or attempt == MAX_RETRIES: |
||||
raise HTTPException(status_code=502, detail=f"图片 URL 下载失败: HTTP {e.code}") from e |
||||
except (urllib.error.URLError, TimeoutError) as e: |
||||
if attempt == MAX_RETRIES: |
||||
raise HTTPException(status_code=504, detail="图片 URL 下载超时或网络连接失败") from e |
||||
time.sleep(2 ** (attempt - 1)) |
||||
|
||||
|
||||
def generate_images( |
||||
*, |
||||
prompt: str, |
||||
size: str, |
||||
ratio: str | None, |
||||
images: list[str] | None, |
||||
profile_name: str, |
||||
): |
||||
"""生成图片并保存到本地，返回结果字典。支持文生图 / 图生图 / 多图合成。""" |
||||
api_key = os.getenv("AGNES_API_KEY") |
||||
if not api_key: |
||||
raise HTTPException(status_code=500, detail="未配置环境变量 AGNES_API_KEY") |
||||
|
||||
resolved_size = _resolve_size(size, ratio) |
||||
payload = _build_payload(prompt=prompt, size=size, ratio=ratio, images=images) |
||||
headers = { |
||||
"Authorization": f"Bearer {api_key}", |
||||
"Content-Type": "application/json", |
||||
} |
||||
response = _post_with_retry(f"{API_BASE_URL}/images/generations", headers, payload) |
||||
|
||||
profile_dir = OUTPUT_DIR / profile_name |
||||
profile_dir.mkdir(parents=True, exist_ok=True) |
||||
saved = [] |
||||
for item in response.get("data") or []: |
||||
try: |
||||
if item.get("b64_json"): |
||||
raw = base64.b64decode(item["b64_json"], validate=True) |
||||
elif item.get("url"): |
||||
raw = _download_image(item["url"]) |
||||
else: |
||||
raise HTTPException(status_code=502, detail="Agnes 图片服务未返回 b64_json 或 url") |
||||
except ValueError as e: |
||||
raise HTTPException(status_code=502, detail="Agnes Base64 图片数据格式无效") from e |
||||
|
||||
filename = f"{DEFAULT_MODEL}-{uuid4().hex}.png" |
||||
(profile_dir / filename).write_bytes(raw) |
||||
saved.append({"filename": filename, "url": f"/output/{profile_name}/{filename}", "bytes": len(raw)}) |
||||
|
||||
return {"model": DEFAULT_MODEL, "size": resolved_size, "ratio": ratio or "1:1", "images": saved} |
||||
Loading…
Reference in new issue