import base64
from fastapi import APIRouter, HTTPException
from fastapi import Body
from huggingface_hub import InferenceClient, TextToImageTargetSize
from io import BytesIO

routerImg = APIRouter()


@routerImg.post("/generate-image/")
async def generate_image(token: str = Body(...), prompt: str = Body(...)):
    try:
        # 创建InferenceClient
        client = InferenceClient("stabilityai/stable-diffusion-3-medium-diffusers", token=token)
        # 使用text_to_image方法生成图片
        image = client.text_to_image(prompt=prompt)

        # 将PIL.Image对象转换为byte数据
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        image_bytes = buffered.getvalue()

        # 将图片转换为base64编码
        image_base64 = base64.b64encode(image_bytes).decode("utf-8")

        # 返回base64编码的图片
        return {"image_base64": image_base64}

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error generating image: {str(e)}")