Add text-to-image route to the application

- Added new textToImage router
- Imported the router from routes
- Included the router in the FastAPI app

app.py

@@ -1,6 +1,6 @@
 from huggingface_hub import InferenceClient
 from fastapi import FastAPI
-from routes import chatCompletion
+from routes import chatCompletion, textToImage
 
 app = FastAPI()
 
@@ -8,4 +8,5 @@ app = FastAPI()
 async def root():
     return {"message": "Hello World"}
 
-app.include_router(chatCompletion.router)
+app.include_router(chatCompletion.router)
+app.include_router(textToImage.router)

models/text_to_image.py

@@ -0,0 +1,15 @@
+from pydantic import BaseModel, Field
+from huggingface_hub import TextToImageTargetSize
+from typing import List, Optional
+
+class TextToImageRequest(BaseModel):
+    prompt: str = Field(..., description="The prompt to generate an image from.")
+    negative_prompt: Optional[List[str]] = Field(None, description="One or several prompts to guide what NOT to include in image generation.")
+    height: Optional[float] = Field(None, description="The height in pixels of the image to generate.")
+    width: Optional[float] = Field(None, description="The width in pixels of the image to generate.")
+    num_inference_steps: Optional[int] = Field(None, description="The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference.")
+    guidance_scale: Optional[float] = Field(None, description="A higher guidance scale value encourages the model to generate images closely linked to the text prompt, but values too high may cause saturation and other artifacts.")
+    model: Optional[str] = Field(None, description="The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed Inference Endpoint. If not provided, the default recommended text-to-image model will be used.")
+    scheduler: Optional[str] = Field(None, description="Override the scheduler with a compatible one.")
+    target_size: Optional[TextToImageTargetSize] = Field(None, description="The size in pixel of the output image")
+    seed: Optional[int] = Field(None, description="Seed for the random number generator.")

routes/textToImage.py

@@ -0,0 +1,25 @@
+from fastapi import APIRouter, Response
+from models.text_to_image import TextToImageRequest
+from huggingface_hub import InferenceClient
+import io
+
+router = APIRouter()
+
+@router.post("/vi/images/generations", tags=["Text to Image"])
+async def text_to_image(body: TextToImageRequest):
+    client = InferenceClient(model=body.model)
+    res = client.text_to_image(
+        prompt=body.prompt,
+        negative_prompt=body.negative_prompt,
+        height=body.height,
+        width=body.width,
+        num_inference_steps=body.num_inference_steps,
+        guidance_scale=body.guidance_scale,
+        scheduler=body.scheduler,
+        target_size=body.target_size,
+        seed=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")