from fastapi import APIRouter, Response
from models.text_to_image import TextToImageRequest
from huggingface_hub import InferenceClient
import io

router = APIRouter()

@router.post("/v1/images/generations", tags=["Text to Image"])
async def text_to_image(t2i_body: TextToImageRequest):
    client = InferenceClient(model=t2i_body.model)
    res = client.text_to_image(
        t2i_body.prompt,
        negative_prompt=t2i_body.negative_prompt,
        height=t2i_body.height,
        width=t2i_body.width,
        num_inference_steps=t2i_body.num_inference_steps,
        guidance_scale=t2i_body.guidance_scale,
        scheduler=t2i_body.scheduler,
       # target_size=t2i_body.target_size,
        seed=t2i_body.seed
    )
    img_byte_arr = io.BytesIO()
    res.save(img_byte_arr, format="PNG")
    img_byte_arr.seek(0)
    try:
        res = client.text_to_image(
            prompt=t2i_body.prompt,
            negative_prompt=t2i_body.negative_prompt,
            height=t2i_body.height,
            width=t2i_body.width,
            num_inference_steps=t2i_body.num_inference_steps,
            guidance_scale=t2i_body.guidance_scale,
            scheduler=t2i_body.scheduler,
            seed=t2i_body.seed
        )
        img_byte_arr = io.BytesIO()
        res.save(img_byte_arr, format="PNG")
        img_byte_arr.seek(0)
        return Response(content=img_byte_arr.getvalue(), media_type="image/png")
    except Exception as e:
        print(f"Error generating image: {e}")
        return {"error": str(e)}, 500