feat: add the 'num_inference_steps' to T2I endpoint (#80)

* add num_inference_steps to text-to-image endpoints

added the generated runner code

final changes for num_inference_steps

removed debug logging

fixed spacing

fixed spacing

* refactor(runner): handle 'num_inference_steps' edge cases

Ensure the pipeline does not crash when the `num_inference_steps`
argument is set to a value lower than 1. Update the default
`num_inference_steps` to align with diffusion model defaults.

---------

Co-authored-by: Rick Staa <rick.staa@outlook.com>

runner/app/pipelines/text_to_image.py

@@ -142,16 +142,16 @@ def __call__(self, prompt: str, **kwargs) -> List[PIL.Image]:
                     torch.Generator(get_torch_device()).manual_seed(s) for s in seed
                 ]
 
+        if "num_inference_steps" in kwargs and kwargs["num_inference_steps"] < 1:
+            del kwargs["num_inference_steps"]
+
         if (
             self.model_id == "stabilityai/sdxl-turbo"
             or self.model_id == "stabilityai/sd-turbo"
         ):
             # SD turbo models were trained without guidance_scale so
             # it should be set to 0
             kwargs["guidance_scale"] = 0.0
-
-            if "num_inference_steps" not in kwargs:
-                kwargs["num_inference_steps"] = 1
         elif SDXL_LIGHTNING_MODEL_ID in self.model_id:
             # SDXL-Lightning models should have guidance_scale = 0 and use
             # the correct number of inference steps for the unet checkpoint loaded

runner/app/routes/text_to_image.py

@@ -25,6 +25,7 @@ class TextToImageParams(BaseModel):
     negative_prompt: str = ""
     seed: int = None
     num_images_per_prompt: int = 1
+    num_inference_steps: int = 50  # TODO: Make optional.
 
 
 responses = {400: {"model": HTTPError}, 500: {"model": HTTPError}}

runner/openapi.json

@@ -438,6 +438,11 @@
                         "type": "integer",
                         "title": "Num Images Per Prompt",
                         "default": 1
+                    },
+                    "num_inference_steps": {
+                        "type": "integer",
+                        "title": "Num Inference Steps",
+                        "default": 50
                     }
                 },
                 "type": "object",

runner/openapi.yaml

@@ -293,6 +293,10 @@ components:
           type: integer
           title: Num Images Per Prompt
           default: 1
+        num_inference_steps:
+          type: integer
+          title: Num Inference Steps
+          default: 50
       type: object
       required:
       - prompt