Tool_Agent/tools/local/flux/main.py

"""BFL FLUX 本地封装 — 异步生图（提交 + 轮询）。

环境变量：
  BFL_API_KEY   必填，请求头 x-key
  BFL_API_BASE  可选，默认 https://api.bfl.ai/v1（全球端点；也可用 api.eu.bfl.ai/v1 等）

接口：
  GET  /health
  POST /submit  提交生成；body.model 为端点路径段，如 flux-2-pro-preview
  POST /query   轮询；polling_url、request_id 来自提交响应

文档: https://docs.bfl.ai/quick_start/generating_images
"""

from __future__ import annotations

import argparse
from typing import Any

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

from bfl_client import poll_result, submit_generation

app = FastAPI(title="BFL FLUX API Proxy")


class SubmitRequest(BaseModel):
    model: str = Field(
        ...,
        description="模型端点路径段，如 flux-2-pro-preview、flux-2-max、flux-dev（对应 /v1/{model}）",
    )
    prompt: str = Field(..., description="文生图提示词")
    width: int | None = Field(default=None, description="输出宽度（像素）")
    height: int | None = Field(default=None, description="输出高度（像素）")
    parameters: dict[str, Any] | None = Field(
        default=None,
        description="合并进请求体的额外字段（官方各模型可选参数）",
    )


class QueryRequest(BaseModel):
    polling_url: str = Field(..., description="提交响应中的 polling_url，须原样使用")
    request_id: str = Field(..., description="提交响应中的 id")


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


@app.post("/submit")
def submit(req: SubmitRequest) -> dict[str, Any]:
    try:
        return submit_generation(
            model=req.model,
            prompt=req.prompt,
            width=req.width,
            height=req.height,
            parameters=req.parameters,
        )
    except ValueError as e:
        raise HTTPException(status_code=503, detail=str(e)) from e
    except RuntimeError as e:
        raise HTTPException(status_code=502, detail=str(e)) from e
    except Exception as e:
        raise HTTPException(status_code=502, detail=str(e)) from e


@app.post("/query")
def query(req: QueryRequest) -> dict[str, Any]:
    try:
        return poll_result(polling_url=req.polling_url, request_id=req.request_id)
    except ValueError as e:
        raise HTTPException(status_code=503, detail=str(e)) from e
    except RuntimeError as e:
        raise HTTPException(status_code=502, detail=str(e)) from e
    except Exception as e:
        raise HTTPException(status_code=502, detail=str(e)) from e


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