perf: pre-warm ImageToVideo SFAST pipeline (#57)

This commit ensures that the model is pre-traced when SFAST is enabled
for the ImageToVideo pipeline. Without this pre-tracing the first
request will be slower than a non SFAST call.

runner/app/pipelines/image_to_image.py

@@ -92,15 +92,25 @@ def __init__(self, model_id: str):
  model_id, **kwargs
  ).to(torch_device)
 
- if os.environ.get("SFAST"):
+ if os.getenv("SFAST", "").strip().lower() == "true":
  logger.info(
- "ImageToImagePipeline will be dynamicallly compiled with stable-fast for %s",
+ "ImageToImagePipeline will be dynamically compiled with stable-fast "
+ "for %s",
  model_id,
  )
  from app.pipelines.sfast import compile_model
 
  self.ldm = compile_model(self.ldm)
 
+ # Warm-up the pipeline.
+ # TODO: Not yet supported for ImageToImagePipeline.
+ if os.getenv("SFAST_WARMUP", "true").lower() == "true":
+ logger.warning(
+ "The 'SFAST_WARMUP' flag is not yet supported for the "
+ "ImageToImagePipeline and will be ignored. As a result the first "
+ "call may be slow if 'SFAST' is enabled."
+ )
+
  def __call__(self, prompt: str, image: PIL.Image, **kwargs) -> List[PIL.Image]:
  seed = kwargs.pop("seed", None)
  if seed is not None:

runner/app/pipelines/image_to_video.py

@@ -8,13 +8,16 @@
 from typing import List
 import logging
 import os
+import time
 
 from PIL import ImageFile
 
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 
 logger = logging.getLogger(__name__)
 
+SFAST_WARMUP_ITERATIONS = 2 # Model warm-up iterations when SFAST is enabled.
+
 
 class ImageToVideoPipeline(Pipeline):
  def __init__(self, model_id: str):
@@ -40,17 +43,52 @@ def __init__(self, model_id: str):
  self.ldm = StableVideoDiffusionPipeline.from_pretrained(model_id, **kwargs)
  self.ldm.to(get_torch_device())
 
- if os.environ.get("SFAST"):
+ if os.getenv("SFAST", "").strip().lower() == "true":
  logger.info(
- "ImageToVideoPipeline will be dynamicallly compiled with stable-fast for %s",
+ "ImageToVideoPipeline will be dynamically compiled with stable-fast "
+ "for %s",
  model_id,
  )
  from app.pipelines.sfast import compile_model
 
  self.ldm = compile_model(self.ldm)
 
+ # Warm-up the pipeline.
+ # NOTE: Initial calls may be slow due to compilation. Subsequent calls will
+ # be faster.
+ if os.getenv("SFAST_WARMUP", "true").lower() == "true":
+ # Retrieve default model params.
+ # TODO: Retrieve defaults from Pydantic class in route.
+ warmup_kwargs = {
+ "image": PIL.Image.new("RGB", (576, 1024)),
+ "height": 576,
+ "width": 1024,
+ "fps": 6,
+ "motion_bucket_id": 127,
+ "noise_aug_strength": 0.02,
+ "decode_chunk_size": 25,
+ }
+
+ logger.info("Warming up ImageToVideoPipeline pipeline...")
+ total_time = 0
+ for ii in range(SFAST_WARMUP_ITERATIONS):
+ t = time.time()
+ try:
+ self.ldm(**warmup_kwargs).frames
+ except Exception as e:
+ # FIXME: When out of memory, pipeline is corrupted.
+ logger.error(f"ImageToVideoPipeline warmup error: {e}")
+ raise e
+ iteration_time = time.time() - t
+ total_time += iteration_time
+ logger.info(
+ "Warmup iteration %s took %s seconds", ii + 1, iteration_time
+ )
+ logger.info("Total warmup time: %s seconds", total_time)
+
  def __call__(self, image: PIL.Image, **kwargs) -> List[List[PIL.Image]]:
  if "decode_chunk_size" not in kwargs:
+ # Decrease decode_chunk_size to reduce memory usage.
  kwargs["decode_chunk_size"] = 4
 
  seed = kwargs.pop("seed", None)

runner/app/pipelines/text_to_image.py

@@ -111,15 +111,25 @@ def __init__(self, model_id: str):
  self.ldm.vae.decode, mode="max-autotune", fullgraph=True
  )
 
- if os.environ.get("SFAST"):
+ if os.getenv("SFAST", "").strip().lower() == "true":
  logger.info(
- "TextToImagePipeline will be dynamicallly compiled with stable-fast for %s",
+ "TextToImagePipeline will be dynamically compiled with stable-fast for "
+ "%s",
  model_id,
  )
  from app.pipelines.sfast import compile_model
 
  self.ldm = compile_model(self.ldm)
 
+ # Warm-up the pipeline.
+ # TODO: Not yet supported for ImageToImagePipeline.
+ if os.getenv("SFAST_WARMUP", "true").lower() == "true":
+ logger.warning(
+ "The 'SFAST_WARMUP' flag is not yet supported for the "
+ "TextToImagePipeline and will be ignored. As a result the first "
+ "call may be slow if 'SFAST' is enabled."
+ )
+
  def __call__(self, prompt: str, **kwargs) -> List[PIL.Image]:
  seed = kwargs.pop("seed", None)
  if seed is not None: