streamdiffusion: Upgrade to tensorrt 10.12 and cuda 10.6 (#646)

* Fix params for streamdfiffusion models

* update streamdiffusion version and move install order

* Update to latest commit on ryan/feat/cnet-spike-trt

* Fix tensorrt==10.9 installation

* Upgrade to CUDA 12.6 and tensorrt 10.12

* Update streamdiffusion version

Also use a tag 😎

---------

Co-authored-by: Victor Elias <[email protected]>

Author: eliteprox

runner/app/live/StreamDiffusionWrapper/wrapper.py

@@ -579,7 +579,7 @@ def create_prefix(
                     unet_model = UNet(
                         fp16=True,
                         device=str(stream.device),
-                        max_batch_size=stream.trt_unet_batch_size,
+                        max_batch=stream.trt_unet_batch_size,
                         min_batch_size=stream.trt_unet_batch_size,
                         embedding_dim=stream.text_encoder.config.hidden_size,
                         unet_dim=stream.unet.config.in_channels,
@@ -602,7 +602,7 @@ def create_prefix(
                     stream.vae.forward = stream.vae.decode
                     vae_decoder_model = VAE(
                         device=str(stream.device),
-                        max_batch_size=self.batch_size
+                        max_batch=self.batch_size
                         if self.mode == "txt2img"
                         else stream.frame_bff_size,
                         min_batch_size=self.batch_size
@@ -630,7 +630,7 @@ def create_prefix(
                     vae_encoder = TorchVAEEncoder(stream.vae).to(torch.device("cuda"))
                     vae_encoder_model = VAEEncoder(
                         device=str(stream.device),
-                        max_batch_size=self.batch_size
+                        max_batch=self.batch_size
                         if self.mode == "txt2img"
                         else stream.frame_bff_size,
                         min_batch_size=self.batch_size

runner/docker/Dockerfile.live-base

@@ -1,5 +1,5 @@
 # Based on https://github.com/huggingface/api-inference-community/blob/main/docker_images/diffusers/Dockerfile
-FROM nvidia/cuda:12.2.2-cudnn8-devel-ubuntu22.04
+FROM nvidia/cuda:12.6.3-cudnn-devel-ubuntu22.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 ENV NVIDIA_VISIBLE_DEVICES=all

runner/docker/Dockerfile.live-base-streamdiffusion

@@ -2,7 +2,7 @@ ARG BASE_IMAGE=livepeer/ai-runner:live-base
 FROM ${BASE_IMAGE}
 
 # Install Miniconda and create the 'comfystream' environment
-ARG PYTHON_VERSION=3.10
+ARG PYTHON_VERSION=3.11
 ENV CONDA_DIR /workspace/miniconda3
 ENV PATH $CONDA_DIR/bin:$PATH
 
@@ -23,19 +23,27 @@ ENV PIP_PREFER_BINARY=1
 
 # Install StreamDiffusion dependencies into the comfystream environment
 RUN conda run -n comfystream pip install --no-cache-dir \
-    torch==2.1.0 \
-    torchvision==0.16.0 \
-    xformers===0.0.22.post7 \
-    --index-url https://download.pytorch.org/whl/cu121
-
-# Install StreamDiffusion @ b623251 (latest at time of writing) into the comfystream environment
-RUN conda run -n comfystream pip install git+https://github.com/cumulo-autumn/StreamDiffusion.git@b623251#egg=streamdiffusion[tensorrt]
-
-# Install TensorRT extension and ensure compatible dependencies into the comfystream environment
-RUN conda run -n comfystream pip install --no-cache-dir --force-reinstall \
-    huggingface-hub==0.25.2 \
-    transformers==4.38.2
-RUN conda run -n comfystream python -m streamdiffusion.tools.install-tensorrt
+    torch==2.7.0 \
+    torchvision==0.22.0 \
+    torchaudio==2.7.0 \
+    xformers===0.0.30 \
+    --index-url https://download.pytorch.org/whl/cu126
+
+# Install StreamDiffusion @ v0.0.1-cnet (latest with controlnet support) into the comfystream environment
+RUN conda run -n comfystream pip install git+https://github.com/livepeer/[email protected]#egg=streamdiffusion[tensorrt]
+
+# Pin versions of ONNX runtime which are too loose on streamdiffusion setup.py
+RUN conda run -n comfystream pip install --no-cache-dir \
+    onnx==1.18.0 onnxruntime==1.22.0 onnxruntime-gpu==1.22.0
+
+# Install TensorRT bindings and libraries into the comfystream environment
+# TODO: update install-tensorrt script in streamdiffusion and use it instead of this
+RUN conda run -n comfystream pip install --no-cache-dir --upgrade --root-user-action=ignore \
+    tensorrt==10.12.0.36 \
+    tensorrt-cu12-bindings==10.12.0.36 \
+    tensorrt-cu12-libs==10.12.0.36 \
+    polygraphy==0.49.24 \
+    onnx-graphsurgeon==0.5.8
 
 WORKDIR /app
 

runner/requirements.live-ai.txt

@@ -12,4 +12,4 @@ nvidia-ml-py==12.560.30
 pynvml==12.0.0
 opencv-python==4.10.0.84
 av==14.4.0
-peft==0.11.1
+peft==0.15.2