Tool_Agent/tools/local/nano_banana/main.py

"""nano_banana — 本地 HTTP 封装 Gemini 原生图模（REST generateContent）。

环境变量：
  GEMINI_API_KEY      必填，对应文档中的 x-goog-api-key
  GEMINI_IMAGE_MODEL  可选，未在请求体指定 model 时使用，默认 gemini-2.5-flash-image
  GEMINI_API_BASE     可选，默认 https://generativelanguage.googleapis.com/v1beta

接口：
  GET  /health
  POST /generate     文生图 / 图+文生图，字段与 registry input_schema 对齐

文档: https://ai.google.dev/gemini-api/docs/image-generation?hl=zh-cn#rest
"""

from __future__ import annotations

import argparse

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

from gemini_image_client import generate_content

app = FastAPI(title="Nano Banana — Gemini Image (REST)")


class ImageInput(BaseModel):
    """参考图：Base64 或 data URL；字段名与 REST inline_data 对应。"""

    mime_type: str = Field(default="image/png", description="如 image/png、image/jpeg")
    data: str = Field(..., description="图片 Base64，或 data:image/...;base64,...")


class GenerateRequest(BaseModel):
    prompt: str = Field(..., description="主提示词")
    model: str | None = Field(
        default=None,
        description=(
            "模型 ID，如 gemini-2.5-flash-image、gemini-3.1-flash-image-preview；"
            "省略则使用 GEMINI_IMAGE_MODEL / 内置默认"
        ),
    )
    aspect_ratio: str | None = Field(
        default=None,
        description='宽高比，如 "1:1"、"16:9"（见官方文档 imageConfig.aspectRatio）',
    )
    image_size: str | None = Field(
        default=None,
        description='Gemini 3.x 输出分辨率：512、1K、2K、4K（generationConfig.imageConfig.imageSize）',
    )
    response_modalities: list[str] | None = Field(
        default=None,
        description='如 ["TEXT","IMAGE"] 或 ["IMAGE"]；省略则由 API 默认',
    )
    images: list[ImageInput] | None = Field(
        default=None,
        description="可选参考图列表（图生图 / 编辑），对应 REST parts 中的 inline_data",
    )


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


@app.post("/generate")
def generate(req: GenerateRequest) -> dict:
    try:
        imgs = None
        if req.images:
            imgs = [{"mime_type": i.mime_type, "data": i.data} for i in req.images]
        return generate_content(
            prompt=req.prompt,
            model=req.model,
            aspect_ratio=req.aspect_ratio,
            image_size=req.image_size,
            response_modalities=req.response_modalities,
            images=imgs,
        )
    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)