Tool_Agent/tools/local/flux/main.py

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

环境变量：
  APIYI_KEY   必填，请求 api.apiyi.com/v1
  APIYI_BASE_URL  可选，默认 https://api.apiyi.com/v1

接口：
  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_flux import generate_flux_image

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

class GenerateRequest(BaseModel):
    prompt: str = Field(..., description="正向提示词")
    aspect_ratio: str = Field(default="16:9", description="宽高比 (如: 1:1, 16:9, 9:16, 2:3, 3:2, 3:7, 7:3)")
    model: str = Field(default="flux-kontext-pro", description="模型选择: 支持 flux-2-flex, flux-2-klein-4b, flux-2-klein-9b, flux-2-max, flux-2-pro, flux-dev, flux-kontext-max, flux-kontext-pro 等")
    seed: Optional[int] = Field(default=None, description="随机种子，用于结果重现")
    safety_tolerance: int = Field(default=2, description="内容安全级别 (0-6)")
    output_format: str = Field(default="png", description="输出格式 jpeg 或 png")
    prompt_upsampling: bool = Field(default=False, description="是否增强提示词")
    image_url: Optional[str] = Field(default=None, description="上传的参考图 URL（开启编辑图生图）")
    mask_url: Optional[str] = Field(default=None, description="蒙版图 URL")

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:
        url = generate_flux_image(
            prompt=req.prompt,
            aspect_ratio=req.aspect_ratio,
            model=req.model,
            seed=req.seed,
            safety_tolerance=req.safety_tolerance,
            output_format=req.output_format,
            prompt_upsampling=req.prompt_upsampling,
            image_url=req.image_url,
            mask_url=req.mask_url
        )
        if url is None:
            raise HTTPException(status_code=500, detail="生成失败，未返回图像链接")
        
        return {"images": [url], "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=8002)
    args = parser.parse_args()
    uvicorn.run(app, host="0.0.0.0", port=args.port)