Update Dockerfile

Dockerfile

@@ -1,4 +1,3 @@
-# Use the base image with your launcher
 FROM thecooltechguy/comfyui_launcher AS launcher
 
 # Create a non-root user first
@@ -27,15 +26,16 @@ RUN chmod +x /start.sh
 RUN mkdir /data && chown user:user /data
 RUN ln -s /data /app
 
-# Install CUDA Toolkit 12.4 base installer and NVIDIA driver
-RUN apt-get update && apt-get install -y wget gnupg2 software-properties-common
+# Install required tools and headers
+RUN apt-get update && apt-get install -y wget gnupg2 software-properties-common build-essential linux-headers-$(uname -r)
+
+# Add NVIDIA's package repositories
 RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.1-1_all.deb
 RUN dpkg -i cuda-keyring_1.1-1_all.deb
 RUN apt-get update
-# Installing CUDA Toolkit 12.4
-RUN apt-get -y install cuda-toolkit-12-4
-# Installing NVIDIA Drivers (Legacy Kernel Module Flavor)
-RUN apt-get install -y cuda-drivers
+
+# Install CUDA Toolkit and NVIDIA drivers
+RUN apt-get -y install cuda-toolkit-12-4 cuda-drivers
 
 # Setting up environment variables for CUDA
 ENV PATH=/usr/local/cuda-12.4/bin:${PATH}
@@ -43,4 +43,4 @@ ENV LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64:${LD_LIBRARY_PATH}
 
 # Switch back to the non-root user for running applications
 USER user
-ENTRYPOINT ["/start.sh"]
+ENTRYPOINT ["/start.sh"]