WIP

runner/app/live/streamer/streamer.py

@@ -124,8 +124,8 @@ async def run_ingress_loop(self, done: Event):
 
                 # crop the max square from the center of the image and scale to 512x512
                 # most models expect this size especially when using tensorrt
-                if frame.width != frame.height:
-                    width, height = frame.size
+                width, height = frame.size
+                if width != height:
                     square_size = min(width, height)
                     frame = frame.crop((width // 2 - square_size // 2, height // 2 - square_size // 2, width // 2 + square_size // 2, height // 2 + square_size // 2))
                 if frame.size != (512, 512):

runner/dl_checkpoints.sh

@@ -79,7 +79,15 @@ function download_all_models() {
 
     # Download live-video-to-video models.
     huggingface-cli download KBlueLeaf/kohaku-v2.1 --include "*.safetensors" "*.json" "*.txt" --exclude ".onnx" ".onnx_data" --cache-dir models
-    huggingface-cli download KwaiVGI/LivePortrait --include "*.safetensors" "*.json" "*.txt" --exclude ".onnx" ".onnx_data" --cache-dir models
+    huggingface-cli download warmshao/FasterLivePortrait --local-dir models/FasterLivePortrait
+}
+
+function build_tensorrt_models() {
+    printf "\nBuilding TensorRT models...\n"
+
+    docker run --rm -v ./models:/models --gpus all \
+        livepeer/ai-runner:live-base-liveportrait \
+        bash -c "cd /app/app/live/FasterLivePortrait && sh scripts/all_onnx2trt.sh && sh scripts/all_onnx2trt_animal.sh"
 }
 
 # Download models with a restrictive license.
@@ -112,7 +120,11 @@ do
         --restricted)
             MODE="restricted"
             shift
-        ;;  
+        ;;
+        --tensorrt)
+            MODE="tensorrt"
+            shift
+        ;;
         --help)
             display_help
             exit 0
@@ -138,6 +150,8 @@ if [ "$MODE" = "beta" ]; then
     download_beta_models
 elif [ "$MODE" = "restricted" ]; then
     download_restricted_models
+elif [ "$MODE" = "tensorrt" ]; then
+    build_tensorrt_models
 else
     download_all_models
 fi

runner/docker/Dockerfile.live-base

@@ -2,6 +2,8 @@
 FROM nvidia/cuda:12.2.2-cudnn8-devel-ubuntu22.04
 
 ENV DEBIAN_FRONTEND=noninteractive
+ENV NVIDIA_VISIBLE_DEVICES=all
+ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility,video
 
 # Install prerequisites
 RUN apt-get update && \

runner/docker/Dockerfile.live-base-liveportrait

@@ -0,0 +1,52 @@
+ARG BASE_IMAGE=livepeer/ai-runner:live-base
+FROM ${BASE_IMAGE}
+
+# Download and install the NVIDIA TensorRT repository local deb
+RUN wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/secure/8.6.1/local_repos/nv-tensorrt-local-repo-ubuntu2204-8.6.1-cuda-12.0_1.0-1_amd64.deb && \
+    dpkg -i nv-tensorrt-local-repo-ubuntu2204-8.6.1-cuda-12.0_1.0-1_amd64.deb && \
+    cp /var/nv-tensorrt-local-repo-ubuntu2204-8.6.1-cuda-12.0/*-keyring.gpg /usr/share/keyrings/ && \
+    rm nv-tensorrt-local-repo-ubuntu2204-8.6.1-cuda-12.0_1.0-1_amd64.deb
+
+# Install TensorRT Python library
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends tensorrt && \
+    rm -rf /var/lib/apt/lists/*
+
+# Install cmake and build dependencies for grid-sample3d-trt-plugin
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends cmake build-essential && \
+    rm -rf /var/lib/apt/lists/*
+
+# Build grid-sample3d-trt-plugin for FasterLivePortrait
+RUN git clone https://github.com/SeanWangJS/grid-sample3d-trt-plugin.git /opt/grid-sample3d-trt-plugin && \
+    cd /opt/grid-sample3d-trt-plugin && \
+    sed -i 's/set_target_properties(${PROJECT_NAME} PROPERTIES CUDA_ARCHITECTURES ".*")/set_target_properties(${PROJECT_NAME} PROPERTIES CUDA_ARCHITECTURES "60;70;75;80;86")/' CMakeLists.txt && \
+    mkdir build && cd build && \
+    export PATH=/usr/local/cuda/bin:$PATH && \
+    cmake .. -DTensorRT_ROOT=/usr/include && \
+    make
+
+# Install required Python version
+ARG PYTHON_VERSION=3.10
+RUN pyenv install $PYTHON_VERSION && \
+    pyenv global $PYTHON_VERSION && \
+    pyenv rehash
+
+# Upgrade pip and install required packages
+ARG PIP_VERSION=23.3.2
+ENV PIP_PREFER_BINARY=1
+RUN pip install --no-cache-dir --upgrade pip==${PIP_VERSION} setuptools==69.5.1 wheel==0.43.0
+
+# Install Python TensorRT packages
+RUN pip install --no-cache-dir tensorrt==8.6.1 numpy==1.26.4
+
+# Clone the FasterLivePortrait repository
+RUN mkdir -p /app/app/live && git clone https://github.com/warmshao/FasterLivePortrait.git /app/app/live/FasterLivePortrait
+WORKDIR /app/app/live/FasterLivePortrait
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Create a symlink on ./checkpoints to the directory that is mounted with the models
+# TODO: Figure out if we can have only 1 of these symlinks
+RUN ln -s /models/FasterLivePortrait ./checkpoints
+RUN ln -s /models/FasterLivePortrait /app/app/checkpoints
+RUN ln -s /models/FasterLivePortrait /app/app/live/checkpoints