diff --git a/build.json b/build.json
index 8fa7bc04..06f2311c 100644
--- a/build.json
+++ b/build.json
@@ -15,6 +15,9 @@
     "AIIDALAB_VERSION": {
       "default": "24.09.0"
     },
+    "AWB_VERSION": {
+      "default": "2.3.0a2"
+    },
     "AIIDALAB_HOME_VERSION": {
       "default": "24.09.0"
     }
diff --git a/docker-bake.hcl b/docker-bake.hcl
index a5e3e697..27d2c6c5 100644
--- a/docker-bake.hcl
+++ b/docker-bake.hcl
@@ -17,6 +17,9 @@ variable "AIIDA_VERSION" {
 variable "AIIDALAB_VERSION" {
 }
 
+variable "AWB_VERSION" {
+}
+
 variable "AIIDALAB_HOME_VERSION" {
 }
 
@@ -88,6 +91,7 @@ target "lab" {
   platforms = "${PLATFORMS}"
   args = {
     "AIIDALAB_VERSION"      = "${AIIDALAB_VERSION}"
+    "AWB_VERSION"      = "${AWB_VERSION}"
     "AIIDALAB_HOME_VERSION" = "${AIIDALAB_HOME_VERSION}"
     "PYTHON_MINOR_VERSION" = get_python_minor_version("${PYTHON_VERSION}")
   }
diff --git a/stack/base/Dockerfile b/stack/base/Dockerfile
index 0cd4c2d4..b5e6a8a1 100644
--- a/stack/base/Dockerfile
+++ b/stack/base/Dockerfile
@@ -5,19 +5,18 @@ LABEL maintainer="AiiDAlab Team <aiidalab@materialscloud.org>"
 
 USER root
 
-# build-essential: includes GCC compilers that are needed when building
-# pip packages from sources, which often seems to happen for pymatgen:
-# https://pymatgen.org/installation.html#installation-tips-for-optional-libraries
 # rsync: needed to support the new AiiDA backup command
 # povray: rendering engine used in aiidalab-widgets-base
-ENV EXTRA_APT_PACKAGES "curl povray rsync build-essential"
+ENV EXTRA_APT_PACKAGES "curl povray rsync"
 
 # For ARM64 we need to install erlang as it is not available on conda-forge
 # (this is needed later as rabbitmq dependency in base-with-services image,
 # but we install it here so that we don't have to invoke apt multiple times.
+# We also install build-essential, which includes GCC compilers that are needed when building
+# pip packages from sources, which is more often needed on ARM64
 ARG TARGETARCH
 RUN if [ "$TARGETARCH" = "arm64" ]; then \
-        EXTRA_APT_PACKAGES="erlang libhdf5-serial-dev pkg-config ${EXTRA_APT_PACKAGES}"; \
+        EXTRA_APT_PACKAGES="erlang build-essential libhdf5-serial-dev pkg-config ${EXTRA_APT_PACKAGES}"; \
     fi;\
     apt-get update --yes && \
     apt-get install --yes --no-install-recommends ${EXTRA_APT_PACKAGES} && \
@@ -51,7 +50,7 @@ COPY pip.conf /etc/pip.conf
 # Install aiida-core and other shared requirements.
 RUN mamba update -y pip mamba zstandard async_generator certipy && \
      mamba install --yes \
-     aiida-core==${AIIDA_VERSION} \
+     aiida-core.atomic_tools==${AIIDA_VERSION} \
      mamba-bash-completion \
      && mamba clean --all -f -y && \
      fix-permissions "${CONDA_DIR}"
diff --git a/stack/lab/Dockerfile b/stack/lab/Dockerfile
index 0c2d29cd..e5e28eae 100644
--- a/stack/lab/Dockerfile
+++ b/stack/lab/Dockerfile
@@ -13,10 +13,11 @@ ENV DOCKER_STACKS_JUPYTER_CMD=notebook
 USER root
 WORKDIR /opt/
 
-# Install aiidalab package
 ARG AIIDALAB_VERSION
+ARG AWB_VERSION
 RUN mamba install --yes \
      aiidalab=${AIIDALAB_VERSION} \
+     aiidalab-widgets-base=${AWB_VERSION} \
      && mamba clean --all -f -y && \
      fix-permissions "${CONDA_DIR}" && \
      fix-permissions "/home/${NB_USER}"