Tool_Agent/tools/local/seedream/main.py

"""APIYI SEEDREAM 本地封装 — 同步生图。

环境变量：
  SEEDREAM_API_KEY / APIYI_KEY 
  APIYI_BASE_URL (可选)

接口：
  GET  /health
  POST /generate
"""

from __future__ import annotations

import argparse
from typing import Any, List, Optional

import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

from generate_seedream import generate_seedream_image

app = FastAPI(title="SeeDream Image Generation APIyi Proxy")


class GenerateRequest(BaseModel):
    prompt: str = Field(..., description="正向提示词")
    model: str = Field(default="seedream-5-0-260128", description="使用的模型名称 (如 seedream-5-0-260128)")
    size: str = Field(default="2048x2048", description="图片尺寸, 如 '1K', '2K', '4K', '2048x2048'")
    quality: str = Field(default="hd", description="图片质量: standard 或 hd")
    n: int = Field(default=1, description="生成数量")
    image_url: Optional[str] = Field(default=None, description="参考图网络外链（开启图生图）")


class GenerateResponse(BaseModel):
    images: List[str]
    status: str


@app.get("/health")
def health() -> dict[str, str]:
    return {"status": "ok"}


@app.post("/generate", response_model=GenerateResponse)
def generate(req: GenerateRequest) -> dict[str, Any]:
    try:
        result = generate_seedream_image(
            prompt=req.prompt,
            model=req.model,
            size=req.size,
            quality=req.quality,
            n=req.n,
            response_format="url",
            image_url=req.image_url
        )
        
        if "error" in result:
            raise HTTPException(status_code=500, detail=result["error"])
            
        if "data" not in result or not result["data"]:
            raise HTTPException(status_code=500, detail="生图失败，返回数据中无图像资料")
            
        urls = [item["url"] for item in result["data"] if "url" in item]
        return {"images": urls, "status": "success"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=8003)
    args = parser.parse_args()
    uvicorn.run(app, host="0.0.0.0", port=args.port)