diff --git a/runner/app/live/streamer/streamer.py b/runner/app/live/streamer/streamer.py index 91d853ac..ae18ff99 100644 --- a/runner/app/live/streamer/streamer.py +++ b/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): diff --git a/runner/dl_checkpoints.sh b/runner/dl_checkpoints.sh index 1b6c64f6..321f13c8 100644 --- a/runner/dl_checkpoints.sh +++ b/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 diff --git a/runner/docker/Dockerfile.live-base b/runner/docker/Dockerfile.live-base index cd94fa7f..6fa161f5 100644 --- a/runner/docker/Dockerfile.live-base +++ b/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 && \ diff --git a/runner/docker/Dockerfile.live-base-liveportrait b/runner/docker/Dockerfile.live-base-liveportrait new file mode 100644 index 00000000..d09335e3 --- /dev/null +++ b/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