diff --git a/Dockerfile b/Dockerfile
index c78c0c2d..087c984c 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.1.0-base-ubuntu20.04 AS base
+FROM nvidia/cuda:12.6.3-cudnn-runtime-ubuntu24.04 AS base
 ENV DEBIAN_FRONTEND=noninteractive
 
 # Install necessary dependencies and Python 3.12
@@ -7,9 +7,6 @@ RUN apt-get update \
     git \
     software-properties-common \
     curl \
-    && add-apt-repository ppa:deadsnakes/ppa \
-    && apt-get update \
-    && apt-get install -y \
     python3.12 \
     python3.12-dev \
     python3.12-venv \
@@ -22,40 +19,27 @@ RUN apt-get update \
     && apt-get clean \
     && rm -rf /var/lib/apt/lists/*
 
-# Install pip for Python 3.12 explicitly
-RUN curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py \
-    && python3.12 get-pip.py \
-    && rm get-pip.py
-
-# Set Python 3.12 as the default python3 and pip
-RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.12 1 \
-    && update-alternatives --install /usr/bin/pip pip /usr/local/bin/pip3 1
-
 # Set the working directory
 WORKDIR /app
 
-# Clone the ComfyUI repository into the working directory
-RUN git clone https://github.com/comfyanonymous/ComfyUI.git /app/comfyui
+# Clone the ComfyUI repository and set up virtual environment
+RUN git clone https://github.com/comfyanonymous/ComfyUI.git /app/comfyui \
+    && python3.12 -m venv /app/venv \
+    && /app/venv/bin/pip install --upgrade pip \
+    && /app/venv/bin/pip install pyyaml \
+    && /app/venv/bin/pip install -r /app/comfyui/requirements.txt
 
-# Install backend dependencies
-RUN python3 -m venv /app/venv
-RUN . venv/bin/activate && pip install --upgrade pip
-RUN . venv/bin/activate && pip install pyyaml
-RUN . venv/bin/activate && pip install -r /app/comfyui/requirements.txt
+# Install PyTorch with CUDA 12.6 support (stable version)
+RUN /app/venv/bin/pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu126
 
-# Install PyTorch with CUDA 12.1 support
-RUN . venv/bin/activate && pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu121
-
-# Clone the ComfyUI-manager repository into a temporary directory, move it, and clean up
-RUN git clone https://github.com/ltdrdata/ComfyUI-Manager.git /app/temp/ComfyUI-Manager && \
-    mv /app/temp/* /app/comfyui/custom_nodes/ && \
-    rm -rf /app/temp
-
-# Install ComfyUI-manager dependencies
-RUN . venv/bin/activate && pip install -r /app/comfyui/custom_nodes/ComfyUI-Manager/requirements.txt
+# Clone ComfyUI-Manager and install its dependencies
+RUN git clone https://github.com/ltdrdata/ComfyUI-Manager.git /app/temp/ComfyUI-Manager \
+    && mv /app/temp/* /app/comfyui/custom_nodes/ \
+    && rm -rf /app/temp \
+    && /app/venv/bin/pip install -r /app/comfyui/custom_nodes/ComfyUI-Manager/requirements.txt
 
 # Expose the backend port
 EXPOSE 8188
 
-# Set the entrypoint command to activate the virtual environment and run the script
-CMD ["/bin/bash", "-c", "source /app/venv/bin/activate && python3 /app/comfyui/main.py --listen 0.0.0.0 --port 8188 && echo 'Server started and ready to accept requests'"]
+# Set the entrypoint to run the app
+CMD ["/bin/bash", "-c", "source /app/venv/bin/activate && python3 /app/comfyui/main.py --listen 0.0.0.0 --port 8188"]
diff --git a/docker-build.sh b/docker-build.sh
index 5c20d985..c3ca5579 100644
--- a/docker-build.sh
+++ b/docker-build.sh
@@ -17,7 +17,7 @@ else
     exit 1
 fi
 
-# Step 2: Start the container without mounting the venv volume
+# Step 2: Start the container without mounting the volumes (venv, custom_nodes) 
 echo "Starting the container..."
 COMPOSE_BAKE=true docker-compose up -d
 if [ $? -eq 0 ]; then
@@ -36,14 +36,14 @@ LOGS_PID=$!  # Save the PID of the background process
 # Wait for the container logs to indicate it's ready (looking for the custom message)
 echo "Waiting for the container to be fully started..."
 while ! docker logs "$container_name" 2>&1 | grep -q "To see the GUI go to: http://0.0.0.0:8188"; do
-    sleep 20
+    sleep 10
 done
 
 # Stop streaming logs (kill the background process)
 kill $LOGS_PID
 echo "Container is fully started."
 
-# Step 4: Copy the 'venv' directory from the container to the host
+# Step 4.1: Copy the 'venv' directory from the container to the host
 echo "Checking if /app/venv exists in the container..."
 if docker exec "$container_name" ls /app/venv; then
     echo "Copying the virtual environment from the container to the host..."
@@ -58,6 +58,21 @@ else
     exit 1
 fi
 
+# Step 4.2: Copy the 'ComfyUI-Manager' directory from the container to the host
+echo "Checking if /app/comfyui/custom_nodes/ComfyUI-Manager exists in the container..."
+if docker exec "$container_name" ls /app/comfyui/custom_nodes/ComfyUI-Manager; then
+    echo "Copying the ComfyUI-Manager from the container to the host..."
+    if ! docker cp "$container_name:/app/comfyui/custom_nodes/ComfyUI-Manager" ./custom_nodes/ComfyUI-Manager; then
+        echo "Failed to copy the ComfyUI-Manager. Exiting."
+        exit 1
+    else
+        echo "ComfyUI-Manager copied successfully."
+    fi
+else
+    echo "/app/comfyui/custom_nodes/ComfyUI-Manager does not exist in the container. Exiting."
+    exit 1
+fi
+
 # Step 5: Stop the container
 echo "Stopping the container..."
 docker-compose down
@@ -68,24 +83,26 @@ else
     exit 1
 fi
 
-# Step 6: Update the Docker Compose file to mount the venv volume
+# Step 6.1: Update the Docker Compose file to mount the venv volume
 echo "Updating Docker Compose file to mount the virtual environment..."
 sed -i '/# Mount the venv directory for persistence/a \ \ \ \ \ \ - ./venv:/app/venv' docker-compose.yml
 if [ $? -eq 0 ]; then
-    echo "Docker Compose file updated successfully."
+    echo "Docker Compose file updated to include venv."
 else
     echo "Failed to update Docker Compose file. Exiting."
     exit 1
 fi
 
-# Step 7: Restart the container with the venv volume mounted
-echo "Restarting the container with the virtual environment mounted..."
-docker-compose up -d
+# Step 6.2: Update the Docker Compose file to mount the venv volume
+echo "Updating Docker Compose file to mount the custom_nodes..."
+sed -i '/# Mount the custom nodes directory directly inside/a \ \ \ \ \ \ - ./custom_nodes:/app/comfyui/custom_nodes' docker-compose.yml
 if [ $? -eq 0 ]; then
-    echo "Container restarted successfully."
+    echo "Docker Compose file updated to include custom_nodes."
 else
-    echo "Failed to restart the container. Exiting."
+    echo "Failed to update Docker Compose file. Exiting."
     exit 1
 fi
 
-echo "Setup complete! The container is running with the virtual environment persisted at ./venv."
\ No newline at end of file
+echo "======================================== SETUP COMPLETE ========================================"
+echo "use 'docker-compose up' to start the container and 'docker-compose down' to stop the container."
+echo "================================================================================================"
\ No newline at end of file
diff --git a/docker-compose.yml b/docker-compose.yml
index be888196..fe46b96a 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -8,14 +8,15 @@ services:
     ports:
       - "8188:8188"  # Expose the backend API or server
     volumes:
-      - ./input:/app/comfyui/input  # Mount the input directory directly inside /app/comfyui
-      - ./models:/app/comfyui/models  # Mount the models directory directly inside /app/comfyui
-      - ./output:/app/comfyui/output  # Mount the output directory directly inside /app/comfyui
-      - ./custom_nodes:/app/comfyui/custom_nodes  # Mount the custom nodes directory directly inside /app/comfyui
-      - ./user:/app/comfyui/user  # Comment this line and uncomment the next line if you want to have a workflows directory in the root directory.
-      #- ./workflows:/app/comfyui/user/default/workflows 
+      - ./input:/app/comfyui/input  # Volume for input files
+      - ./models:/app/comfyui/models  # Volume for models
+      - ./output:/app/comfyui/output  # Volume for output files
+      - ./user:/app/comfyui/user  # Volume for user settings, Also this is where your workflows are stored
+      #- ./workflows:/app/comfyui/user/default/workflows  # Uncomment if you to make a workflow directory in the comfyui directory
 
-      # Mount the venv directory for persistence(automatically mounted with you run docker-build.sh)
+      # Mount the venv directory for persistence (automatically mounted with you run docker-build.sh) #Don't change this
+
+      # Mount the custom nodes directory directly inside /app/comfyui (automatically mounted with you run docker-build.sh)  #Don't change this
 
 
     environment: