dl_checkpoints.sh - restore docker pull for mutable docker tags (#375)

livepeer · Dec 18, 2024 · 18aab2c · 18aab2c
1 parent 5ef92f6
commit 18aab2c
Showing 1 changed file with 15 additions and 3 deletions.
diff --git a/runner/dl_checkpoints.sh b/runner/dl_checkpoints.sh
@@ -128,8 +128,12 @@ function build_tensorrt_models() {
     # StreamDiffusion (compile a matrix of models and timesteps)
     MODELS="stabilityai/sd-turbo KBlueLeaf/kohaku-v2.1"
     TIMESTEPS="3 4" # This is basically the supported sizes for the t_index_list
+    AI_RUNNER_STREAMDIFFUSION_IMAGE=${AI_RUNNER_STREAMDIFFUSION_IMAGE:-livepeer/ai-runner:live-app-streamdiffusion}
+    docker pull $AI_RUNNER_STREAMDIFFUSION_IMAGE
+    # ai-worker has tags hardcoded in `var livePipelineToImage` so we need to use the same tag in here:
+    docker image tag $AI_RUNNER_STREAMDIFFUSION_IMAGE livepeer/ai-runner:live-app-streamdiffusion
     docker run --rm -v ./models:/models --gpus all -l TensorRT-engines \
-        ${AI_RUNNER_STREAMDIFFUSION_IMAGE:-livepeer/ai-runner:live-app-streamdiffusion} \
+        $AI_RUNNER_STREAMDIFFUSION_IMAGE \
         bash -c "for model in $MODELS; do
                     for timestep in $TIMESTEPS; do
                         echo \"Building TensorRT engines for model=\$model timestep=\$timestep...\" && \
@@ -138,8 +142,12 @@ function build_tensorrt_models() {
                 done"
 
     # FasterLivePortrait
+    AI_RUNNER_LIVEPORTRAIT_IMAGE=${AI_RUNNER_LIVEPORTRAIT_IMAGE:-livepeer/ai-runner:live-app-liveportrait}
+    docker pull $AI_RUNNER_LIVEPORTRAIT_IMAGE
+    # ai-worker has tags hardcoded in `var livePipelineToImage` so we need to use the same tag in here:
+    docker image tag $AI_RUNNER_LIVEPORTRAIT_IMAGE livepeer/ai-runner:live-app-liveportrait
     docker run --rm -v ./models:/models --gpus all -l TensorRT-engines  \
-        ${AI_RUNNER_LIVEPORTRAIT_IMAGE:-livepeer/ai-runner:live-app-liveportrait} \
+        $AI_RUNNER_LIVEPORTRAIT_IMAGE \
         bash -c "cd /app/app/live/FasterLivePortrait && \
                     if [ ! -f '/models/FasterLivePortrait--checkpoints/liveportrait_onnx/stitching_lip.trt' ]; then
                         echo 'Building TensorRT engines for LivePortrait models (regular)...'
@@ -155,8 +163,12 @@ function build_tensorrt_models() {
                     fi"
 
     # ComfyUI (only DepthAnything for now)
+    AI_RUNNER_COMFYUI_IMAGE=${AI_RUNNER_COMFYUI_IMAGE:-livepeer/ai-runner:live-app-comfyui}
+    docker pull $AI_RUNNER_COMFYUI_IMAGE
+    # ai-worker has tags hardcoded in `var livePipelineToImage` so we need to use the same tag in here:
+    docker image tag $AI_RUNNER_COMFYUI_IMAGE livepeer/ai-runner:live-app-comfyui
     docker run --rm -v ./models:/models --gpus all -l TensorRT-engines \
-        ${AI_RUNNER_COMFYUI_IMAGE:-livepeer/ai-runner:live-app-comfyui} \
+        $AI_RUNNER_COMFYUI_IMAGE \
         bash -c "cd /comfyui/models/Depth-Anything-Onnx && \
                     python /comfyui/custom_nodes/ComfyUI-Depth-Anything-Tensorrt/export_trt.py && \
                     mkdir -p /comfyui/models/tensorrt/depth-anything && \