diff --git a/build_and_push_docker b/build_and_push_docker
new file mode 100644
index 00000000..e891e22c
--- /dev/null
+++ b/build_and_push_docker
@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+
+# Extract version from version.py
+if [ -f "version.py" ]; then
+    VERSION=$(python3 -c "from version import VERSION; print(VERSION)")
+    echo "Building version: $VERSION"
+else
+    echo "Error: version.py not found. Please create a version.py file with VERSION defined."
+    exit 1
+fi
+
+echo "Docker builds from the repo, not this dir. Make sure changes are pushed to the repo."
+echo "Building version: $VERSION and latest"
+# wait 2 seconds
+sleep 2
+
+# Build the image with cache busting
+docker build --build-arg CACHEBUST=$(date +%s) -t aitoolkit:$VERSION -f docker/Dockerfile .
+
+# Tag with version and latest
+docker tag aitoolkit:$VERSION ostris/aitoolkit:$VERSION
+docker tag aitoolkit:$VERSION ostris/aitoolkit:latest
+
+# Push both tags
+echo "Pushing images to Docker Hub..."
+docker push ostris/aitoolkit:$VERSION
+docker push ostris/aitoolkit:latest
+
+echo "Successfully built and pushed ostris/aitoolkit:$VERSION and ostris/aitoolkit:latest"
\ No newline at end of file
diff --git a/build_and_push_docker.yaml b/build_and_push_docker.yaml
deleted file mode 100644
index bbda385a..00000000
--- a/build_and_push_docker.yaml
+++ /dev/null
@@ -1,8 +0,0 @@
-#!/usr/bin/env bash
-
-echo "Docker builds from the repo, not this dir. Make sure changes are pushed to the repo."
-# wait 2 seconds
-sleep 2
-docker build --build-arg CACHEBUST=$(date +%s) -t aitoolkit:latest -f docker/Dockerfile .
-docker tag aitoolkit:latest ostris/aitoolkit:latest
-docker push ostris/aitoolkit:latest
\ No newline at end of file
diff --git a/docker-compose.yml b/docker-compose.yml
new file mode 100644
index 00000000..a589feed
--- /dev/null
+++ b/docker-compose.yml
@@ -0,0 +1,24 @@
+version: "3.8"
+
+services:
+  ai-toolkit:
+    image: ostris/aitoolkit:latest
+    restart: unless-stopped
+    ports:
+      - "8675:8675"
+    volumes:
+      - ~/.cache/huggingface/hub:/root/.cache/huggingface/hub
+      - ./aitk_db.db:/app/ai-toolkit/aitk_db.db
+      - ./datasets:/app/ai-toolkit/datasets
+      - ./output:/app/ai-toolkit/output
+    environment:
+      - AI_TOOLKIT_AUTH=${AI_TOOLKIT_AUTH:-password}
+      - NODE_ENV=production
+      - TZ=UTC
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
diff --git a/docker/Dockerfile b/docker/Dockerfile
index e98572c4..e07b781e 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -1,31 +1,66 @@
-FROM runpod/base:0.6.2-cuda12.2.0
+FROM nvidia/cuda:12.6.3-base-ubuntu22.04
 
 LABEL authors="jaret"
 
+# Set noninteractive to avoid timezone prompts
+ENV DEBIAN_FRONTEND=noninteractive
+
 # Install dependencies
-RUN apt-get update
+RUN apt-get update && apt-get install --no-install-recommends -y \
+    git \
+    curl \
+    build-essential \
+    cmake \
+    wget \
+    python3.10 \
+    python3-pip \
+    python3-dev \
+    python3-setuptools \
+    python3-wheel \
+    python3-venv \
+    ffmpeg \
+    tmux \
+    htop \
+    nvtop \
+    python3-opencv \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install nodejs
+WORKDIR /tmp
+RUN curl -sL https://deb.nodesource.com/setup_23.x -o nodesource_setup.sh && \
+    bash nodesource_setup.sh && \
+    apt-get update && \
+    apt-get install -y nodejs && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
 
 WORKDIR /app
-ARG CACHEBUST=1
-RUN git clone https://github.com/ostris/ai-toolkit.git && \
+
+# Fix cache busting by moving CACHEBUST to right before git clone
+ARG CACHEBUST=1234
+RUN echo "Cache bust: ${CACHEBUST}" && \
+    git clone https://github.com/ostris/ai-toolkit.git && \
     cd ai-toolkit && \
     git submodule update --init --recursive
 
 WORKDIR /app/ai-toolkit
 
-RUN ln -s /usr/bin/python3 /usr/bin/python
-RUN python -m pip install -r requirements.txt
+# Set aliases for python and pip
+RUN echo 'alias python=python3' >> /etc/bash.bashrc && \
+    echo 'alias pip=pip3' >> /etc/bash.bashrc
 
-RUN apt-get install -y tmux nvtop htop
+# Install Python dependencies
+RUN pip install --no-cache-dir torch==2.6.0 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cu126 && \
+    pip install --no-cache-dir -r requirements.txt
 
-RUN pip install jupyterlab
+# Build UI
+WORKDIR /app/ai-toolkit/ui
+RUN npm install && \
+    npm run build && \
+    npm run update_db
 
-# mask workspace
-RUN mkdir /workspace
+# Expose port (assuming the application runs on port 3000)
+EXPOSE 8675
 
-
-# symlink app to workspace
-RUN ln -s /app/ai-toolkit /workspace/ai-toolkit
-
-WORKDIR /
-CMD ["/start.sh"]
\ No newline at end of file
+CMD ["npm", "run", "start"]
\ No newline at end of file
diff --git a/version.py b/version.py
new file mode 100644
index 00000000..d0076416
--- /dev/null
+++ b/version.py
@@ -0,0 +1 @@
+VERSION = "0.2.0"
\ No newline at end of file