diff --git a/runner/Dockerfile b/runner/Dockerfile
index d9ed5cd9..10bb9deb 100644
--- a/runner/Dockerfile
+++ b/runner/Dockerfile
@@ -18,14 +18,14 @@ ENV PYENV_ROOT /root/.pyenv
 ENV PATH $PYENV_ROOT/shims:$PYENV_ROOT/bin:$PATH
 
 # Install your desired Python version
-ARG PYTHON_VERSION=3.11
+ARG PYTHON_VERSION=3.10
 RUN pyenv install $PYTHON_VERSION && \
     pyenv global $PYTHON_VERSION && \
     pyenv rehash
 
 # Upgrade pip and install your desired packages
 ARG PIP_VERSION=23.3.2
-RUN pip install --no-cache-dir --upgrade pip==${PIP_VERSION} setuptools==69.5.1 wheel==0.43.0 
+RUN pip install --no-cache-dir --upgrade pip==${PIP_VERSION} setuptools==69.5.1 wheel==0.43.0
 
 # Set environment variables
 ENV MAX_WORKERS=1
@@ -41,4 +41,4 @@ ENV MODEL_DIR=/models
 # - Setting the init system to runit
 
 # We're keeping the WORKDIR /app here as it's a good default
-WORKDIR /app
\ No newline at end of file
+WORKDIR /app
diff --git a/runner/docker/Dockerfile.stream-diffusion b/runner/docker/Dockerfile.stream-diffusion
index 1fc3dd66..0d10d304 100644
--- a/runner/docker/Dockerfile.stream-diffusion
+++ b/runner/docker/Dockerfile.stream-diffusion
@@ -20,9 +20,10 @@ RUN git clone https://github.com/cumulo-autumn/StreamDiffusion.git /streamdiffus
 WORKDIR /streamdiffusion
 
 # Install StreamDiffusion
-RUN pip install streamdiffusion[tensorrt]
+RUN pip install git+https://github.com/cumulo-autumn/StreamDiffusion.git@main#egg=streamdiffusion[tensorrt]
 
 # Install TensorRT extension
+RUN python -c "import platform; print(platform.system())"
 RUN python -m streamdiffusion.tools.install-tensorrt
 
 # Set environment variables for NVIDIA drivers