Sfast optimization for T2I , I2I and Upscale models

.github/workflows/ai-runner-docker.yaml

@@ -76,4 +76,4 @@ jobs:
  file: "Dockerfile"
  labels: ${{ steps.meta.outputs.labels }}
  cache-from: type=registry,ref=livepeerci/build:cache
- cache-to: type=registry,ref=livepeerci/build:cache,mode=max
+ cache-to: type=registry,ref=livepeerci/build:cache,mode=max

.github/workflows/validate-openapi-on-pr.yaml

@@ -46,4 +46,4 @@ jobs:
  if ! git diff --exit-code; then
  echo "::error::Go bindings have changed. Please run 'make' at the root of the repository and commit the changes."
  exit 1
- fi
+ fi

go.mod

@@ -5,7 +5,7 @@ go 1.21.5
 require (
  github.com/deepmap/oapi-codegen/v2 v2.2.0
  github.com/docker/cli v24.0.5+incompatible
- github.com/docker/docker v24.0.7+incompatible
+ github.com/docker/docker v24.0.9+incompatible
  github.com/docker/go-connections v0.4.0
  github.com/getkin/kin-openapi v0.124.0
  github.com/go-chi/chi/v5 v5.0.12

go.sum

@@ -17,8 +17,8 @@ github.com/docker/cli v24.0.5+incompatible h1:WeBimjvS0eKdH4Ygx+ihVq1Q++xg36M/rM
 github.com/docker/cli v24.0.5+incompatible/go.mod h1:JLrzqnKDaYBop7H2jaqPtU4hHvMKP+vjCwu2uszcLI8=
 github.com/docker/distribution v2.8.3+incompatible h1:AtKxIZ36LoNK51+Z6RpzLpddBirtxJnzDrHLEKxTAYk=
 github.com/docker/distribution v2.8.3+incompatible/go.mod h1:J2gT2udsDAN96Uj4KfcMRqY0/ypR+oyYUYmja8H+y+w=
-github.com/docker/docker v24.0.7+incompatible h1:Wo6l37AuwP3JaMnZa226lzVXGA3F9Ig1seQen0cKYlM=
-github.com/docker/docker v24.0.7+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
+github.com/docker/docker v24.0.9+incompatible h1:HPGzNmwfLZWdxHqK9/II92pyi1EpYKsAqcl4G0Of9v0=
+github.com/docker/docker v24.0.9+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
 github.com/docker/go-connections v0.4.0 h1:El9xVISelRB7BuFusrZozjnkIM5YnzCViNKohAFqRJQ=
 github.com/docker/go-connections v0.4.0/go.mod h1:Gbd7IOopHjR8Iph03tsViu4nIes5XhDvyHbTtUxmeec=
 github.com/docker/go-units v0.4.0 h1:3uh0PgVws3nIA0Q+MwDC8yjEPf9zjRfZZWXZYDct3Tw=

runner/app/main.py

@@ -1,5 +1,7 @@
 import logging
 import os
+import sys
+import cv2
 from contextlib import asynccontextmanager
 
 from app.routes import health
@@ -15,8 +17,8 @@ async def lifespan(app: FastAPI):
 
  app.include_router(health.router)
 
- pipeline = os.environ["PIPELINE"]
- model_id = os.environ["MODEL_ID"]
+ pipeline = os.environ.get("PIPELINE", "") # Default to 
+ model_id = os.environ.get("MODEL_ID", "") # Provide a default if necessary
 
  app.pipeline = load_pipeline(pipeline, model_id)
  app.include_router(load_route(pipeline))
@@ -44,8 +46,10 @@ def load_pipeline(pipeline: str, model_id: str) -> any:
  from app.pipelines.audio_to_text import AudioToTextPipeline
 
  return AudioToTextPipeline(model_id)
- case "frame-interpolation":
- raise NotImplementedError("frame-interpolation pipeline not implemented")
+ case "FILMPipeline":
+ from app.pipelines.frame_interpolation import FILMPipeline
+
+ return FILMPipeline(model_id)
  case "upscale":
  from app.pipelines.upscale import UpscalePipeline
 
@@ -78,8 +82,10 @@ def load_route(pipeline: str) -> any:
  from app.routes import audio_to_text
 
  return audio_to_text.router
- case "frame-interpolation":
- raise NotImplementedError("frame-interpolation pipeline not implemented")
+ case "FILMPipeline":
+ from app.routes import frame_interpolation
+
+ return frame_interpolation.router
  case "upscale":
  from app.routes import upscale
 

runner/app/pipelines/frame_interpolation.py

@@ -1,5 +1,113 @@
-from app.pipelines.base import Pipeline
+import torch
+from torchvision.transforms import v2
+from tqdm import tqdm
+import bisect
+import numpy as np
+from app.pipelines.utils.utils import get_model_dir
 
 
-class FrameInterpolationPipeline(Pipeline):
- pass
+class FILMPipeline:
+ model: torch.jit.ScriptModule
+
+ def __init__(self, model_id: str):
+ self.model_id = model_id
+ model_dir = get_model_dir() # Get the directory where models are stored
+ model_path = f"{model_dir}/{model_id}" # Construct the full path to the model file
+
+ self.model = torch.jit.load(model_path, map_location="cpu")
+ self.model.eval()
+
+ def to(self, *args, **kwargs):
+ self.model = self.model.to(*args, **kwargs)
+ return self
+
+ @property
+ def device(self) -> torch.device:
+ # Checking device for ScriptModule requires checking one of its parameters
+ params = self.model.parameters()
+ return next(params).device
+
+ @property
+ def dtype(self) -> torch.dtype:
+ # Checking device for ScriptModule requires checking one of its parameters
+ params = self.model.parameters()
+ return next(params).dtype
+
+ def __call__(
+ self,
+ reader,
+ writer,
+ inter_frames: int = 2,
+ ):
+ transforms = v2.Compose(
+ [
+ v2.ToDtype(torch.uint8, scale=True),
+ ]
+ )
+
+ writer.open()
+
+ while True:
+ frame_1 = reader.get_frame()
+ # If the first frame read is None then there are no more frames
+ if frame_1 is None:
+ break
+
+ frame_2 = reader.get_frame()
+ # If the second frame read is None there there is a final frame
+ if frame_2 is None:
+ writer.write_frame(transforms(frame_1))
+ break
+
+ # frame_1 and frame_2 must be tensors with n c h w format
+ frame_1 = frame_1.unsqueeze(0)
+ frame_2 = frame_2.unsqueeze(0)
+
+ frames = inference(
+ self.model, frame_1, frame_2, inter_frames, self.device, self.dtype
+ )
+
+ frames = [transforms(frame.detach().cpu()) for frame in frames]
+ for frame in frames:
+ writer.write_frame(frame)
+
+ writer.close()
+
+
+def inference(
+ model, img_batch_1, img_batch_2, inter_frames, device, dtype
+) -> torch.Tensor:
+ results = [img_batch_1, img_batch_2]
+
+ idxes = [0, inter_frames + 1]
+ remains = list(range(1, inter_frames + 1))
+
+ splits = torch.linspace(0, 1, inter_frames + 2)
+
+ for _ in tqdm(range(len(remains)), "Generating in-between frames"):
+ starts = splits[idxes[:-1]]
+ ends = splits[idxes[1:]]
+ distances = (
+ (splits[None, remains] - starts[:, None])
+ / (ends[:, None] - starts[:, None])
+ - 0.5
+ ).abs()
+ matrix = torch.argmin(distances).item()
+ start_i, step = np.unravel_index(matrix, distances.shape)
+ end_i = start_i + 1
+
+ x0 = results[start_i].to(device=device, dtype=dtype)
+ x1 = results[end_i].to(device=device, dtype=dtype)
+
+ dt = x0.new_full((1, 1), (splits[remains[step]] - splits[idxes[start_i]])) / (
+ splits[idxes[end_i]] - splits[idxes[start_i]]
+ )
+
+ with torch.no_grad():
+ prediction = model(x0, x1, dt)
+ insert_position = bisect.bisect_left(idxes, remains[step])
+ idxes.insert(insert_position, remains[step])
+ results.insert(insert_position, prediction.clamp(0, 1).float())
+ del remains[step]
+
+ return results

runner/app/pipelines/image_to_image.py

@@ -1,6 +1,7 @@
 import logging
 import os
 from enum import Enum
+import time
 from typing import List, Optional, Tuple
 
 import PIL
@@ -29,6 +30,7 @@
 
 logger = logging.getLogger(__name__)
 
+SFAST_WARMUP_ITERATIONS = 2 # Model warm-up iterations when SFAST is enabled.
 
 class ModelName(Enum):
  """Enumeration mapping model names to their corresponding IDs."""
@@ -146,11 +148,29 @@ def __init__(self, model_id: str):
  # 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."
- )
+ warmup_kwargs = {
+ "prompt":"A warmed up pipeline is a happy pipeline a short poem by ricksta",
+ "image": PIL.Image.new("RGB", (576, 1024)),
+ "strength": 0.8,
+ "negative_prompt": "No blurry or weird artifacts",
+ "num_images_per_prompt":4,
+ }
+
+ logger.info("Warming up ImageToImagePipeline pipeline...")
+ total_time = 0
+ for ii in range(SFAST_WARMUP_ITERATIONS):
+ t = time.time()
+ try:
+ self.ldm(**warmup_kwargs).images
+ except Exception as e:
+ logger.error(f"ImageToImagePipeline 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)
 
  if deepcache_enabled and not (
  is_lightning_model(model_id) or is_turbo_model(model_id)

runner/app/pipelines/image_to_video.py

@@ -6,10 +6,14 @@
 import PIL
 import torch
 from app.pipelines.base import Pipeline
-from app.pipelines.utils import SafetyChecker, get_model_dir, get_torch_device
 from diffusers import StableVideoDiffusionPipeline
 from huggingface_hub import file_download
 from PIL import ImageFile
+from app.pipelines.utils import (
+ SafetyChecker,
+ get_model_dir,
+ get_torch_device
+)
 
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 

runner/app/pipelines/text_to_image.py

@@ -1,6 +1,7 @@
 import logging
 import os
 from enum import Enum
+import time
 from typing import List, Optional, Tuple
 
 import PIL
@@ -28,6 +29,7 @@
 
 logger = logging.getLogger(__name__)
 
+SFAST_WARMUP_ITERATIONS = 2 # Model warm-up iterations when SFAST is enabled.
 
 class ModelName(Enum):
  """Enumeration mapping model names to their corresponding IDs."""
@@ -173,14 +175,31 @@ def __init__(self, model_id: str):
 
  self.ldm = compile_model(self.ldm)
 
- # Warm-up the pipeline.
- # TODO: Not yet supported for TextToImagePipeline.
  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."
- )
+ # Retrieve default model params.
+ # TODO: Retrieve defaults from Pydantic class in route.
+ warmup_kwargs = {
+ "prompt": "A happy pipe in the line looking at the wall with words sfast",
+ "num_images_per_prompt": 4,
+ "negative_prompt": "No blurry or weird artifacts",
+ }
+
+ logger.info("Warming up TextToImagePipeline pipeline...")
+ total_time = 0
+ for ii in range(SFAST_WARMUP_ITERATIONS):
+ t = time.time()
+ try:
+ self.ldm(**warmup_kwargs).images
+ except Exception as e:
+ # FIXME: When out of memory, pipeline is corrupted.
+ logger.error(f"TextToImagePipeline 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)
 
  if deepcache_enabled and not (
  is_lightning_model(model_id) or is_turbo_model(model_id)

runner/app/pipelines/upscale.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import time
 from typing import List, Optional, Tuple
 
 import PIL
@@ -20,6 +21,7 @@
 
 logger = logging.getLogger(__name__)
 
+SFAST_WARMUP_ITERATIONS = 2 # Model warm-up iterations when SFAST is enabled.
 
 class UpscalePipeline(Pipeline):
  def __init__(self, model_id: str):
@@ -67,11 +69,29 @@ def __init__(self, model_id: str):
  # Warm-up the pipeline.
  # TODO: Not yet supported for UpscalePipeline.
  if os.getenv("SFAST_WARMUP", "true").lower() == "true":
- logger.warning(
- "The 'SFAST_WARMUP' flag is not yet supported for the "
- "UpscalePipeline and will be ignored. As a result the first "
- "call may be slow if 'SFAST' is enabled."
- )
+ # Retrieve default model params.
+ # TODO: Retrieve defaults from Pydantic class in route.
+ warmup_kwargs = {
+ "prompt": "Upscaling the pipeline with sfast enabled",
+ "image": PIL.Image.new("RGB", (400, 400)), # anything higher than this size cause the model to OOM
+ }
+
+ logger.info("Warming up Upscale pipeline...")
+ total_time = 0
+ for ii in range(SFAST_WARMUP_ITERATIONS):
+ t = time.time()
+ try:
+ self.ldm(**warmup_kwargs).images
+ except Exception as e:
+ # FIXME: When out of memory, pipeline is corrupted.
+ logger.error(f"Upscale pipeline 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)
 
  if deepcache_enabled and not (
  is_lightning_model(model_id) or is_turbo_model(model_id)
@@ -86,7 +106,7 @@ def __init__(self, model_id: str):
  elif deepcache_enabled:
  logger.warning(
  "DeepCache is not supported for Lightning or Turbo models. "
- "TextToImagePipeline will NOT be optimized with DeepCache for %s",
+ "UpscalingPiepline will NOT be optimized with DeepCache for %s",
  model_id,
  )
 

runner/app/pipelines/utils/__init__.py

@@ -11,4 +11,8 @@
  is_turbo_model,
  split_prompt,
  validate_torch_device,
+ frames_compactor,
+ video_shredder,
+ DirectoryReader,
+ DirectoryWriter
 )