diff --git a/.github/workflows/build-docker-image.yaml b/.github/workflows/build-docker-image.yaml
new file mode 100644
index 00000000..7f2bd6e3
--- /dev/null
+++ b/.github/workflows/build-docker-image.yaml
@@ -0,0 +1,120 @@
+name: Build Docker Image
+
+on:
+    workflow_call:
+        inputs:
+            build_args:
+                required: false
+                default: ""
+                type: string
+            cache_id:
+                required: true
+                type: string
+            extract_flavor:
+                required: false
+                default: ""
+                type: string
+            image_name:
+                required: true
+                type: string
+            image_tag:
+                required: false
+                default: ""
+                type: string
+            registry:
+                required: false
+                default: ghcr.io
+                type: string
+
+env:
+    FULL_IMAGE_NAME: ${{ inputs.registry }}/${{ inputs.image_name }}
+
+jobs:
+    build-image:
+        runs-on: ubuntu-latest
+        permissions:
+            contents: read
+            packages: write
+        strategy:
+            fail-fast: false
+            matrix:
+                platform:
+                    - linux/amd64
+                    - linux/arm64
+
+        steps:
+            - name: Prepare
+              run: |
+                  platform=${{ matrix.platform }}
+                  echo "PLATFORM_PAIR=${platform//\//-}" >> $GITHUB_ENV
+
+            - name: Checkout repository
+              uses: actions/checkout@v4
+
+            - name: Set up QEMU
+              uses: docker/setup-qemu-action@v3
+
+            - name: Set up Docker Buildx
+              uses: docker/setup-buildx-action@v3
+
+            - name: Log in to the Container registry
+              uses: docker/login-action@v3
+              with:
+                  registry: ${{ inputs.registry }}
+                  username: ${{ github.actor }}
+                  password: ${{ secrets.GITHUB_TOKEN }}
+
+            - name: Extract metadata for Docker images
+              id: meta
+              uses: docker/metadata-action@v5
+              with:
+                  images: ${{ env.FULL_IMAGE_NAME }}
+                  tags: |
+                      type=ref,event=branch
+                      type=ref,event=tag
+                      type=sha,prefix=git-
+                      type=semver,pattern={{version}}
+                      type=semver,pattern={{major}}.{{minor}}
+                      ${{ inputs.image_tag }}
+                  flavor: |
+                      latest=${{ github.ref == 'refs/heads/main' }}
+                      ${{ inputs.extract_flavor }}
+
+            - name: Extract metadata for Docker cache
+              id: cache-meta
+              uses: docker/metadata-action@v5
+              with:
+                  images: ${{ env.FULL_IMAGE_NAME }}
+                  tags: |
+                      type=ref,event=branch
+                  flavor: |
+                      prefix=cache-${{ inputs.cache_id }}-${{ matrix.platform }}-
+
+            - name: Build Docker image
+              uses: docker/build-push-action@v5
+              id: build
+              with:
+                  context: .
+                  push: true
+                  platforms: ${{ matrix.platform }}
+                  labels: ${{ steps.meta.outputs.labels }}
+                  outputs: type=image,name=${{ env.FULL_IMAGE_NAME }},push-by-digest=true,name-canonical=true,push=true
+                  cache-from: type=registry,ref=${{ steps.cache-meta.outputs.tags }}
+                  cache-to: type=registry,ref=${{ steps.cache-meta.outputs.tags }},mode=max
+                  build-args: |
+                      BUILD_HASH=${{ github.sha }}
+                      ${{ inputs.build_args }}
+
+            - name: Export digest
+              run: |
+                  mkdir -p /tmp/digests
+                  digest="${{ steps.build.outputs.digest }}"
+                  touch "/tmp/digests/${digest#sha256:}"
+
+            - name: Upload digest
+              uses: actions/upload-artifact@v4
+              with:
+                  name: digests-${{ inputs.cache_id }}-${{ env.PLATFORM_PAIR }}
+                  path: /tmp/digests/*
+                  if-no-files-found: error
+                  retention-days: 1
diff --git a/.github/workflows/docker-build.yaml b/.github/workflows/docker-build.yaml
index b38634e1..5dab1b59 100644
--- a/.github/workflows/docker-build.yaml
+++ b/.github/workflows/docker-build.yaml
@@ -1,277 +1,41 @@
 name: Create and publish Docker images with specific build args
 
 on:
-  workflow_dispatch:
-  push:
-    branches:
-      - main
-      - dev
-
-env:
-  REGISTRY: ghcr.io
-  IMAGE_NAME: ${{ github.repository }}
-  FULL_IMAGE_NAME: ghcr.io/${{ github.repository }}
+    workflow_dispatch:
+    push:
+        branches:
+            - main
+            - dev
 
 jobs:
-  build-main-image:
-    runs-on: ubuntu-latest
-    permissions:
-      contents: read
-      packages: write
-    strategy:
-      fail-fast: false
-      matrix:
-        platform:
-          - linux/amd64
-          - linux/arm64
-
-    steps:
-      - name: Prepare
-        run: |
-          platform=${{ matrix.platform }}
-          echo "PLATFORM_PAIR=${platform//\//-}" >> $GITHUB_ENV
-
-      - name: Checkout repository
-        uses: actions/checkout@v4
-
-      - name: Set up QEMU
-        uses: docker/setup-qemu-action@v3
-
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
-
-      - name: Log in to the Container registry
-        uses: docker/login-action@v3
-        with:
-          registry: ${{ env.REGISTRY }}
-          username: ${{ github.actor }}
-          password: ${{ secrets.GITHUB_TOKEN }}
-
-      - name: Extract metadata for Docker images (default latest tag)
-        id: meta
-        uses: docker/metadata-action@v5
-        with:
-          images: ${{ env.FULL_IMAGE_NAME }}
-          tags: |
-            type=ref,event=branch
-            type=ref,event=tag
-            type=sha,prefix=git-
-            type=semver,pattern={{version}}
-            type=semver,pattern={{major}}.{{minor}}
-          flavor: |
-            latest=${{ github.ref == 'refs/heads/main' }}
-
-      - name: Extract metadata for Docker cache
-        id: cache-meta
-        uses: docker/metadata-action@v5
-        with:
-          images: ${{ env.FULL_IMAGE_NAME }}
-          tags: |
-            type=ref,event=branch
-          flavor: |
-            prefix=cache-${{ matrix.platform }}-
-
-      - name: Build Docker image (latest)
-        uses: docker/build-push-action@v5
-        id: build
-        with:
-          context: .
-          push: true
-          platforms: ${{ matrix.platform }}
-          labels: ${{ steps.meta.outputs.labels }}
-          outputs: type=image,name=${{ env.FULL_IMAGE_NAME }},push-by-digest=true,name-canonical=true,push=true
-          cache-from: type=registry,ref=${{ steps.cache-meta.outputs.tags }}
-          cache-to: type=registry,ref=${{ steps.cache-meta.outputs.tags }},mode=max
-          build-args: |
-            BUILD_HASH=${{ github.sha }}
-
-      - name: Export digest
-        run: |
-          mkdir -p /tmp/digests
-          digest="${{ steps.build.outputs.digest }}"
-          touch "/tmp/digests/${digest#sha256:}"
-
-      - name: Upload digest
-        uses: actions/upload-artifact@v4
-        with:
-          name: digests-main-${{ env.PLATFORM_PAIR }}
-          path: /tmp/digests/*
-          if-no-files-found: error
-          retention-days: 1
-
-  build-cuda-image:
-    runs-on: ubuntu-latest
-    permissions:
-      contents: read
-      packages: write
-    strategy:
-      fail-fast: false
-      matrix:
-        platform:
-          - linux/amd64
-          - linux/arm64
-
-    steps:
-      - name: Prepare
-        run: |
-          platform=${{ matrix.platform }}
-          echo "PLATFORM_PAIR=${platform//\//-}" >> $GITHUB_ENV
-
-      - name: Checkout repository
-        uses: actions/checkout@v4
-
-      - name: Set up QEMU
-        uses: docker/setup-qemu-action@v3
-
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
-
-      - name: Log in to the Container registry
-        uses: docker/login-action@v3
+    build-main-image:
+        uses: ./.github/workflows/build-docker-image.yaml
         with:
-          registry: ${{ env.REGISTRY }}
-          username: ${{ github.actor }}
-          password: ${{ secrets.GITHUB_TOKEN }}
+            image_name: ${{ github.repository }}
+            cache_id: main
 
-      - name: Extract metadata for Docker images (cuda tag)
-        id: meta
-        uses: docker/metadata-action@v5
+    build-cuda-image:
+        uses: ./.github/workflows/build-docker-image.yaml
         with:
-          images: ${{ env.FULL_IMAGE_NAME }}
-          tags: |
-            type=ref,event=branch
-            type=ref,event=tag
-            type=sha,prefix=git-
-            type=semver,pattern={{version}}
-            type=semver,pattern={{major}}.{{minor}}
-            type=raw,enable=${{ github.ref == 'refs/heads/main' }},prefix=,suffix=,value=cuda
-          flavor: |
-            latest=${{ github.ref == 'refs/heads/main' }}
-            suffix=-cuda,onlatest=true
+            image_name: ${{ github.repository }}
+            cache_id: cuda
+            image_tag: type=raw,enable=${{ github.ref == 'refs/heads/main' }},prefix=,suffix=,value=cuda
+            extract_flavor: suffix=-cuda,onlatest=true
+            build_args: |
+                USE_CUDA=true
 
-      - name: Extract metadata for Docker cache
-        id: cache-meta
-        uses: docker/metadata-action@v5
+    merge-main-images:
+        uses: ./.github/workflows/merge-docker-images.yaml
+        needs: [build-main-image]
         with:
-          images: ${{ env.FULL_IMAGE_NAME }}
-          tags: |
-            type=ref,event=branch
-          flavor: |
-            prefix=cache-cuda-${{ matrix.platform }}-
+            image_name: ${{ github.repository }}
+            cache_id: main
 
-      - name: Build Docker image (cuda)
-        uses: docker/build-push-action@v5
-        id: build
+    merge-cuda-images:
+        uses: ./.github/workflows/merge-docker-images.yaml
+        needs: [build-cuda-image]
         with:
-          context: .
-          push: true
-          platforms: ${{ matrix.platform }}
-          labels: ${{ steps.meta.outputs.labels }}
-          outputs: type=image,name=${{ env.FULL_IMAGE_NAME }},push-by-digest=true,name-canonical=true,push=true
-          cache-from: type=registry,ref=${{ steps.cache-meta.outputs.tags }}
-          cache-to: type=registry,ref=${{ steps.cache-meta.outputs.tags }},mode=max
-          build-args: |
-            BUILD_HASH=${{ github.sha }}
-            USE_CUDA=true
-
-      - name: Export digest
-        run: |
-          mkdir -p /tmp/digests
-          digest="${{ steps.build.outputs.digest }}"
-          touch "/tmp/digests/${digest#sha256:}"
-
-      - name: Upload digest
-        uses: actions/upload-artifact@v4
-        with:
-          name: digests-cuda-${{ env.PLATFORM_PAIR }}
-          path: /tmp/digests/*
-          if-no-files-found: error
-          retention-days: 1
-
-  merge-main-images:
-    runs-on: ubuntu-latest
-    needs: [build-main-image]
-    steps:
-      - name: Download digests
-        uses: actions/download-artifact@v4
-        with:
-          pattern: digests-main-*
-          path: /tmp/digests
-          merge-multiple: true
-
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
-
-      - name: Log in to the Container registry
-        uses: docker/login-action@v3
-        with:
-          registry: ${{ env.REGISTRY }}
-          username: ${{ github.actor }}
-          password: ${{ secrets.GITHUB_TOKEN }}
-
-      - name: Extract metadata for Docker images (default latest tag)
-        id: meta
-        uses: docker/metadata-action@v5
-        with:
-          images: ${{ env.FULL_IMAGE_NAME }}
-          tags: |
-            type=ref,event=branch
-            type=ref,event=tag
-            type=sha,prefix=git-
-          flavor: |
-            latest=${{ github.ref == 'refs/heads/main' }}
-
-      - name: Create manifest list and push
-        working-directory: /tmp/digests
-        run: |
-          docker buildx imagetools create $(jq -cr '.tags | map("-t " + .) | join(" ")' <<< "$DOCKER_METADATA_OUTPUT_JSON") \
-            $(printf '${{ env.FULL_IMAGE_NAME }}@sha256:%s ' *)
-
-      - name: Inspect image
-        run: |
-          docker buildx imagetools inspect ${{ env.FULL_IMAGE_NAME }}:${{ steps.meta.outputs.version }}
-
-  merge-cuda-images:
-    runs-on: ubuntu-latest
-    needs: [build-cuda-image]
-    steps:
-      - name: Download digests
-        uses: actions/download-artifact@v4
-        with:
-          pattern: digests-cuda-*
-          path: /tmp/digests
-          merge-multiple: true
-
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
-
-      - name: Log in to the Container registry
-        uses: docker/login-action@v3
-        with:
-          registry: ${{ env.REGISTRY }}
-          username: ${{ github.actor }}
-          password: ${{ secrets.GITHUB_TOKEN }}
-
-      - name: Extract metadata for Docker images (default latest tag)
-        id: meta
-        uses: docker/metadata-action@v5
-        with:
-          images: ${{ env.FULL_IMAGE_NAME }}
-          tags: |
-            type=ref,event=branch
-            type=ref,event=tag
-            type=sha,prefix=git-
-            type=raw,enable=${{ github.ref == 'refs/heads/main' }},prefix=,suffix=,value=cuda
-          flavor: |
-            latest=${{ github.ref == 'refs/heads/main' }}
-            suffix=-cuda,onlatest=true
-
-      - name: Create manifest list and push
-        working-directory: /tmp/digests
-        run: |
-          docker buildx imagetools create $(jq -cr '.tags | map("-t " + .) | join(" ")' <<< "$DOCKER_METADATA_OUTPUT_JSON") \
-            $(printf '${{ env.FULL_IMAGE_NAME }}@sha256:%s ' *)
-
-      - name: Inspect image
-        run: |
-          docker buildx imagetools inspect ${{ env.FULL_IMAGE_NAME }}:${{ steps.meta.outputs.version }}
+            image_name: ${{ github.repository }}
+            cache_id: cuda
+            extract_flavor: suffix=-cuda,onlatest=true
+            extract_tags: type=raw,enable=${{ github.ref == 'refs/heads/main' }},prefix=,suffix=,value=cuda
diff --git a/.github/workflows/merge-docker-images.yaml b/.github/workflows/merge-docker-images.yaml
new file mode 100644
index 00000000..512f42cf
--- /dev/null
+++ b/.github/workflows/merge-docker-images.yaml
@@ -0,0 +1,71 @@
+name: Merge Docker Images
+
+on:
+    workflow_call:
+        inputs:
+            cache_id:
+                required: true
+                type: string
+            extract_flavor:
+                required: false
+                default: ""
+                type: string
+            extract_tags:
+                required: false
+                default: ""
+                type: string
+            image_name:
+                required: true
+                type: string
+            registry:
+                required: false
+                default: ghcr.io
+                type: string
+
+env:
+    FULL_IMAGE_NAME: ${{ inputs.registry }}/${{ inputs.image_name }}
+
+jobs:
+    merge-images:
+        runs-on: ubuntu-latest
+        steps:
+            - name: Download digests
+              uses: actions/download-artifact@v4
+              with:
+                  pattern: digests-${{ inputs.cache_id }}-*
+                  path: /tmp/digests
+                  merge-multiple: true
+
+            - name: Set up Docker Buildx
+              uses: docker/setup-buildx-action@v3
+
+            - name: Log in to the Container registry
+              uses: docker/login-action@v3
+              with:
+                  registry: ${{ inputs.registry }}
+                  username: ${{ github.actor }}
+                  password: ${{ secrets.GITHUB_TOKEN }}
+
+            - name: Extract metadata for Docker images
+              id: meta
+              uses: docker/metadata-action@v5
+              with:
+                  images: ${{ env.FULL_IMAGE_NAME }}
+                  tags: |
+                      type=ref,event=branch
+                      type=ref,event=tag
+                      type=sha,prefix=git-
+                      ${{ inputs.extract_tags }}
+                  flavor: |
+                      latest=${{ github.ref == 'refs/heads/main' }}
+                      ${{ inputs.extract_flavor }}
+
+            - name: Create manifest list and push
+              working-directory: /tmp/digests
+              run: |
+                  docker buildx imagetools create $(jq -cr '.tags | map("-t " + .) | join(" ")' <<< "$DOCKER_METADATA_OUTPUT_JSON") \
+                    $(printf '${{ env.FULL_IMAGE_NAME }}@sha256:%s ' *)
+
+            - name: Inspect image
+              run: |
+                  docker buildx imagetools inspect ${{ env.FULL_IMAGE_NAME }}:${{ steps.meta.outputs.version }}
diff --git a/Dockerfile b/Dockerfile
index ac5864da..5d2677df 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,4 +1,4 @@
-FROM python:3.11-slim-bookworm as base
+FROM python:3.11-slim-bookworm AS base
 
 # Use args
 ARG USE_CUDA