diff --git a/Dockerfile b/Dockerfile
index db5b2267e..86f98bf95 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,42 +1,33 @@
+# Stage 1: Compile image
 FROM python:3.10-slim-bullseye AS compile-image
 WORKDIR /app
 
-RUN apt-get update &&  \
+RUN apt-get update && \
     apt-get install --no-install-recommends -y wget libpq-dev gcc g++ python3-dev && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
 RUN python -m venv /opt/venv
-# Make sure we use the virtualenv:
 ENV PATH="/opt/venv/bin:$PATH"
 
 COPY requirements.txt .
-RUN pip install --upgrade pip
-RUN pip install --no-cache-dir -r requirements.txt
+RUN pip install --upgrade pip && \
+    pip install --no-cache-dir -r requirements.txt
 
 COPY . .
-COPY entrypoint.sh ./entrypoint.sh
-COPY wait-for-it.sh ./wait-for-it.sh
-RUN chmod +x ./entrypoint.sh ./wait-for-it.sh
 
-# Downloads the tools
-RUN python superagi/tool_manager.py
-
-# Set executable permissions for install_tool_dependencies.sh
-RUN chmod +x install_tool_dependencies.sh
-
-# Install dependencies
-RUN ./install_tool_dependencies.sh
+RUN chmod +x ./entrypoint.sh ./wait-for-it.sh ./install_tool_dependencies.sh ./entrypoint_celery.sh
 
+# Stage 2: Build image
 FROM python:3.10-slim-bullseye AS build-image
 WORKDIR /app
-RUN apt-get update && apt-get install --no-install-recommends -y libpq-dev && \
+
+RUN apt-get update && \
+    apt-get install --no-install-recommends -y libpq-dev && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
 COPY --from=compile-image /opt/venv /opt/venv
 COPY --from=compile-image /app /app
 
-# Make sure we use the virtualenv:
 ENV PATH="/opt/venv/bin:$PATH"
-
diff --git a/docker-compose.yaml b/docker-compose.yaml
index ce9f8134a..35a089433 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -20,7 +20,7 @@ services:
       - super__postgres
     networks:
       - super_network
-    command: ["celery", "-A", "superagi.worker", "worker", "--beat","--loglevel=info"]
+    command: ["/app/entrypoint_celery.sh"]
   gui:
     build:
       context: ./gui
diff --git a/entrypoint.sh b/entrypoint.sh
index 9345436dc..8474e6e4b 100755
--- a/entrypoint.sh
+++ b/entrypoint.sh
@@ -1,5 +1,10 @@
 #!/bin/bash
 
+# Downloads the tools
+python superagi/tool_manager.py
+
+# Install dependencies
+./install_tool_dependencies.sh
 
 # Run Alembic migrations
 alembic upgrade head
diff --git a/entrypoint_celery.sh b/entrypoint_celery.sh
old mode 100644
new mode 100755
index 0b2b8a531..a5d7b5491
--- a/entrypoint_celery.sh
+++ b/entrypoint_celery.sh
@@ -1,5 +1,9 @@
 #!/bin/bash
-Xvfb :0 -screen 0 1280x1024x24 &
-x11vnc -display :0 -N -forever -shared &
 
-exec "$@"
\ No newline at end of file
+# Downloads the tools
+python superagi/tool_manager.py
+
+# Install dependencies
+./install_tool_dependencies.sh
+
+exec celery -A superagi.worker worker --beat --loglevel=info
\ No newline at end of file