| 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,
|
|
|
| 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 |
