Update Dockerfile

Dockerfile

@@ -3,6 +3,10 @@ FROM ubuntu:22.04
 # Set environment variables
 ENV DEBIAN_FRONTEND=noninteractive
 ENV HOME=/root
+# CUDA paths for when running on GPU (RunPod)
+ENV CUDA_HOME=/usr/local/cuda
+ENV PATH=${CUDA_HOME}/bin:${PATH}
+ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
 
 # Install dependencies, code-server, and Ollama
 RUN apt-get update && \
@@ -17,7 +21,7 @@ RUN apt-get update && \
     # Install pip for Python 3.12
     curl -sS https://bootstrap.pypa.io/get-pip.py | python3.12 && \
     # Upgrade pip
-    python3 -m pip install --upgrade pip && \
+    python3 -m pip install --upgrade pip setuptools wheel && \
     # Install Node.js 22.x from NodeSource
     curl -fsSL https://deb.nodesource.com/setup_22.x | bash - && \
     apt-get install -y nodejs && \
@@ -32,6 +36,16 @@ RUN npm install -g @anthropic-ai/claude-code @anthropic-ai/dxt
 # Create a directory for the workspace
 RUN mkdir -p /workspace
 
+# Copy requirements files (if they exist)
+COPY requirements*.txt /workspace/
+
+# Create a requirements-cpu.txt without vllm for build time
+RUN if [ -f /workspace/requirements.txt ]; then \
+        grep -v "vllm" /workspace/requirements.txt | grep -v "xformers==.*+.*" > /workspace/requirements-cpu.txt || true; \
+        echo "xformers==0.0.32" >> /workspace/requirements-cpu.txt; \
+        pip3 install --no-cache-dir -r /workspace/requirements-cpu.txt || true; \
+    fi
+
 # Create configuration directory for code-server and Ollama
 RUN mkdir -p /root/.config/code-server /root/.ollama
 
@@ -46,8 +60,14 @@ RUN code-server --install-extension ms-python.python && \
     code-server --install-extension ritwickdey.LiveServer && \
     code-server --install-extension ms-toolsai.jupyter
 
-# Create a startup script
+# Create a startup script that handles vllm installation on GPU
 RUN echo '#!/bin/bash\n\
+# Check if running on GPU and install vllm if needed\n\
+if [ -d "/usr/local/cuda" ] && ! python3 -c "import vllm" 2>/dev/null; then\n\
+    echo "GPU detected, installing vllm for Blackwell support..."\n\
+    pip3 install git+https://github.com/vllm-project/vllm.git@04e1642e3 || echo "vllm installation failed, continuing..."\n\
+fi\n\
+\n\
 # Start Ollama in the background\n\
 /usr/local/bin/ollama serve &\n\
 \n\