Merge pull request #8 from ssoto/feat/implements_image_generation_improvement

VAE encoder

Author: ssoto

ai_platform/api/public/image_tasks/router.py

@@ -1,7 +1,8 @@
 import logging
 from fastapi.responses import JSONResponse
+from pydantic import Field, BaseModel, field_validator
 from fastapi import APIRouter, Request, status, BackgroundTasks
-from fastapi.encoders import jsonable_encoder
+from typing import Optional
 
 from ai_platform.domain.image_tasks.models import ImageTask
 from ai_platform.domain.image_tasks.use_cases import create_image_task, afind_image_task_by_id
@@ -16,6 +17,28 @@
 )
 
 
+class GenerationRequest(BaseModel):
+    prompt: str = Field(
+        ...,
+        title="Prompt to generate the image"
+    )
+    generation_steps: Optional[int] = Field(
+        50,
+        title="Number of steps to generate the image. Values between 1 and 999"
+    )
+    seed: Optional[int] = Field(
+        None,
+        title="Seed to generate the image"
+    )
+
+    @field_validator("generation_steps") # noqa
+    @classmethod
+    def validate_generation_steps(cls, v):
+        if v < 1 or v > 999:
+            raise ValueError("Generation steps must be between 1 and 999")
+        return v
+
+
 @router.get("/")
 async def retrieve(request: Request, id_task: str):
 
@@ -34,13 +57,13 @@ async def retrieve(request: Request, id_task: str):
 @router.post("/")
 async def generate(
         request: Request,
-        prompt: str,
+        body: GenerationRequest,
         background_tasks: BackgroundTasks
 ):
-    prompt = jsonable_encoder(prompt)
-
     image_task = ImageTask(
-        prompt=prompt,
+        prompt=body.prompt,
+        generation_steps=body.generation_steps,
+        seed=body.seed
     )
     # FIXME: this image service is a local endpoint, it should be a service
     image_task.url = get_image_url(image_task.id)

ai_platform/domain/image_tasks/models.py

@@ -61,7 +61,7 @@ class ImageTask(BaseModel):
         description="Task status"
     )
     reason: Optional[str] = Field(
-        default=None,
+        None,
         title="Reason of failure",
         description="Reason of failure"
     )
@@ -70,6 +70,11 @@ class ImageTask(BaseModel):
         title="Number of steps to generate the image",
         description="Number of steps to generate the image"
     )
+    seed: Optional[int] = Field(
+        None,
+        title="Seed to generate the image",
+        description="Seed to generate the image"
+    )
     url: str = Field(
         None,
         title="URL to download the image",

ai_platform/task_queue/images_creation.py

@@ -2,6 +2,8 @@
 import platform
 
 from diffusers import DiffusionPipeline
+from diffusers.models import AutoencoderKL
+from diffusers import StableDiffusionPipeline
 import torch
 
 from ai_platform.domain.image_tasks.models import ImageTask
@@ -13,12 +15,16 @@
 
 
 def startup_pipeline(only_download=False):
-    # https://huggingface.co/docs/diffusers/tutorials/basic_training
-    pipe = DiffusionPipeline.from_pretrained(
-        "runwayml/stable-diffusion-v1-5",
-        torch_dtype=torch.float16,
-        variant="fp16",
+    logger.info(f"Starting pipeline")
+    logger.info("Loading VAE model")
+    vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")
+    model = "runwayml/stable-diffusion-v1-5"
+    logger.info(f"Loading model: {model}")
+    pipe = StableDiffusionPipeline.from_pretrained(
+        model,
+        vae=vae
     )
+
     if only_download:
         return
 
@@ -40,11 +46,16 @@ def startup_pipeline(only_download=False):
 
 
 def create_image(pipe: DiffusionPipeline, image: ImageTask):
+    kwargs = {}
+    if image.seed is not None:
+        generator = torch.Generator(device="cpu").manual_seed(image.seed)
+        kwargs["generator"] = generator
+
+    kwargs["prompt"] = image.prompt
+    kwargs["num_inference_steps"] = image.generation_steps
+
     # Results match those from the CPU device after the warmup pass.
-    result = pipe(
-        image.prompt,
-        num_inference_steps=image.generation_steps
-    )
+    result = pipe(**kwargs)
     image_file = result.images[0]
     logger.info(f"Image generated: {image_file}")
     return image_file