diff --git a/.dockerignore b/.dockerignore
index c04bc49..6593f2f 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -1 +1,2 @@
 poetry.lock
+docker/
diff --git a/README.md b/README.md
index 66c6a72..b2b7d0a 100644
--- a/README.md
+++ b/README.md
@@ -143,26 +143,22 @@ have them.
 
 ### CPU Image
 ```bash
-docker build --target selfie-cpu -t selfie:cpu .
+./docker/build.sh cpu 
 
 docker run -p 8181:8181 \
   --name selfie-cpu \
   -v $(pwd)/data:/selfie/data \
-  -v $(pwd)/selfie:/selfie/selfie \
-  -v $(pwd)/selfie-ui:/selfie/selfie-ui \
   -v $HOME/.cache/huggingface:/root/.cache/huggingface \
   selfie:cpu
 ```
 
 ### Nvidia GPU Image
 ```bash
-docker build --target selfie-gpu -t selfie:gpu .
+./docker/build.sh gpu 
 
 docker run -p 8181:8181 \
   --name selfie-gpu \
   -v $(pwd)/data:/selfie/data \
-  -v $(pwd)/selfie:/selfie/selfie \
-  -v $(pwd)/selfie-ui:/selfie/selfie-ui \
   -v $HOME/.cache/huggingface:/root/.cache/huggingface \
   --gpus all \
   --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 \
@@ -170,22 +166,29 @@ docker run -p 8181:8181 \
 ````
 
 ### MacOS ARM64 Image
+
+
+> Disclaimer: 
+> 
+> This Docker container is designed to run on a wide range of architectures, including Apple's M1 and M2 chips.
+However, due to current limitations with Docker on macOS, direct access to Metal APIs is not available for containers.
+As a result, applications requiring intensive graphics processing may experience reduced performance compared to running natively on macOS.
+This setup is recommended for development and testing purposes only. Running it on a native machine is recommended for better performance.
+> 
+> We're continuously exploring ways to improve performance and compatibility.
+
+
+
 ```bash
-DOCKER_BUILDKIT=0 docker build --target selfie-arm64 -t selfie:arm64 .
+ ./docker/build.sh arm64
 
 docker run -p 8181:8181 \
   --name selfie-arm64 \
   -v $(pwd)/data:/selfie/data \
-  -v $(pwd)/selfie:/selfie/selfie \
-  -v $(pwd)/selfie-ui:/selfie/selfie-ui \
   -v $HOME/.cache/huggingface:/root/.cache/huggingface \
   selfie:arm64
 ```
 
-> **Note**: on an M1/M2/M3 Mac, you may need to prefix the build command with `DOCKER_BUILDKIT=0` to disable BuildKit.
-> ```bash
-> DOCKER_BUILDKIT=0 docker build -t selfie:arm64 .
-> ```
 
 ## Setting Up Selfie
 
diff --git a/docker/Dockerfile.arm64 b/docker/Dockerfile.arm64
new file mode 100644
index 0000000..cecf8ab
--- /dev/null
+++ b/docker/Dockerfile.arm64
@@ -0,0 +1,20 @@
+FROM node:18-alpine as selfie-ui
+WORKDIR /selfie
+COPY selfie-ui/package.json selfie-ui/yarn.lock ./
+RUN yarn install --frozen-lockfile --non-interactive
+COPY selfie-ui/ .
+RUN yarn run build
+
+FROM python:3.11 as selfie-arm64
+ENV PYTHONDONTWRITEBYTECODE 1
+ENV PYTHONUNBUFFERED 1
+ENV PIP_NO_CACHE_DIR=1
+
+WORKDIR /selfie
+COPY . .
+COPY --from=selfie-ui /selfie/out/ ./selfie-ui/out
+RUN pip install poetry --no-cache-dir
+RUN poetry config virtualenvs.create false
+RUN poetry install --no-interaction --no-ansi
+EXPOSE 8181
+CMD ["python", "-m", "selfie", "--gpu", "--verbose"]
\ No newline at end of file
diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu
new file mode 100644
index 0000000..e32a4c5
--- /dev/null
+++ b/docker/Dockerfile.cpu
@@ -0,0 +1,19 @@
+FROM node:18-alpine as selfie-ui
+WORKDIR /selfie
+COPY selfie-ui/package.json selfie-ui/yarn.lock ./
+RUN yarn install --frozen-lockfile --non-interactive
+COPY selfie-ui/ .
+RUN yarn run build
+
+FROM python:3.11 as selfie-cpu
+ENV PYTHONDONTWRITEBYTECODE 1
+ENV PYTHONUNBUFFERED 1
+ENV PIP_NO_CACHE_DIR=1
+WORKDIR /selfie
+COPY . .
+COPY --from=selfie-ui /selfie/out/ ./selfie-ui/out
+RUN pip install poetry --no-cache-dir
+RUN poetry config virtualenvs.create false
+RUN poetry install --no-interaction --no-ansi
+EXPOSE 8181
+CMD ["python", "-m", "selfie"]
\ No newline at end of file
diff --git a/Dockerfile b/docker/Dockerfile.gpu
similarity index 55%
rename from Dockerfile
rename to docker/Dockerfile.gpu
index b28d447..8630ed6 100644
--- a/Dockerfile
+++ b/docker/Dockerfile.gpu
@@ -1,35 +1,22 @@
-# Base image for building the UI
-FROM node:18.19-alpine3.18 AS selfie-ui
+FROM node:18-alpine as selfie-ui
 WORKDIR /selfie
 COPY selfie-ui/package.json selfie-ui/yarn.lock ./
 RUN yarn install --frozen-lockfile --non-interactive
 COPY selfie-ui/ .
 RUN yarn run build
 
-# Base image for the application
-FROM python:3.11 as selfie-base
-ENV PYTHONDONTWRITEBYTECODE 1
-ENV PYTHONUNBUFFERED 1
-ENV PIP_NO_CACHE_DIR=1
-WORKDIR /selfie
-COPY . .
-COPY --from=selfie-ui /selfie/out/ ./selfie-ui/out
-RUN pip install poetry --no-cache-dir
-RUN poetry config virtualenvs.create false
-RUN poetry install --no-interaction --no-ansi
-EXPOSE 8181
-
-# CPU-specific configuration
-FROM selfie-base as selfie-cpu
-CMD ["python", "-m", "selfie"]
-
 # GPU-specific configuration using an NVIDIA base image
 # Ensure that the chosen image version has wheels available at https://jllllll.github.io/llama-cpp-python-cuBLAS-wheels.
 # You can see what CUDA version an image uses at https://docs.nvidia.com/deeplearning/frameworks/support-matrix/index.html.
 # For example, the image nvcr.io/nvidia/pytorch:23.10-py3 uses CUDA 12.2, which is available at https://jllllll.github.io/llama-cpp-python-cuBLAS-wheels/AVX2/cu122.
 FROM nvcr.io/nvidia/pytorch:23.10-py3 as selfie-gpu
-COPY --from=selfie-base /selfie /selfie
+ENV PYTHONDONTWRITEBYTECODE 1
+ENV PYTHONUNBUFFERED 1
+ENV PIP_NO_CACHE_DIR=1
+
 WORKDIR /selfie
+COPY . .
+COPY --from=selfie-ui /selfie/out/ ./selfie-ui/out
 
 RUN pip install poetry --no-cache-dir
 ENV PATH="/root/.local/bin:$PATH"
@@ -38,16 +25,4 @@ RUN poetry install --no-interaction --no-ansi
 RUN bash /selfie/scripts/llama-cpp-python-cublas.sh
 
 # --verbose and other options should probably be controlled by the user
-CMD ["poetry", "run", "python", "-m", "selfie", "--gpu", "--verbose"]
-
-# ARM64-specific configuration (Apple Silicon)
-FROM selfie-base as selfie-arm64
-# Uncomment below if additional dependencies are needed for ARM64
-# RUN apt-get update && apt-get install -y --no-install-recommends <arm64-specific-dependencies> && rm -rf /var/lib/apt/lists/*
-CMD if [ "$(uname -m)" = "arm64" ]; then \
-        echo "Running with GPU support"; \
-        python -m selfie --gpu; \
-    else \
-        echo "Running without GPU support"; \
-        python -m selfie; \
-    fi
+CMD ["poetry", "run", "python", "-m", "selfie", "--gpu", "--verbose"]
\ No newline at end of file
diff --git a/docker/Dockerfile.ui b/docker/Dockerfile.ui
new file mode 100644
index 0000000..236d386
--- /dev/null
+++ b/docker/Dockerfile.ui
@@ -0,0 +1,6 @@
+FROM node:18-alpine as selfie-ui
+WORKDIR /selfie
+COPY selfie-ui/package.json selfie-ui/yarn.lock ./
+RUN yarn install --frozen-lockfile --non-interactive
+COPY selfie-ui/ .
+RUN yarn run build
\ No newline at end of file
diff --git a/docker/build.sh b/docker/build.sh
new file mode 100755
index 0000000..2b75ce2
--- /dev/null
+++ b/docker/build.sh
@@ -0,0 +1,24 @@
+#!/bin/bash
+
+TARGET_ARCH=$1
+
+if [ -z "$TARGET_ARCH" ]; then
+    echo "Please provide the target architecture (cpu, gpu, or arm64) as an argument."
+    exit 1
+fi
+
+case $TARGET_ARCH in
+    cpu)
+        docker build -t selfie:cpu -f docker/Dockerfile.ui -f docker/Dockerfile.cpu .
+        ;;
+    gpu)
+        docker build -t selfie:gpu -f docker/Dockerfile.ui -f docker/Dockerfile.gpu .
+        ;;
+    arm64)
+        docker build -t selfie:arm64 -f docker/Dockerfile.ui -f docker/Dockerfile.arm64 .
+        ;;
+    *)
+        echo "Invalid target architecture. Please choose from cpu, gpu, or arm64."
+        exit 1
+        ;;
+esac
\ No newline at end of file