diff --git a/Dockerfile b/Dockerfile
index 62f8796..c946fb8 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,6 +1,16 @@
+FROM lsiobase/alpine:3.12 as buildstage
+
+RUN \
+ apk add --no-cache \
+    git && \
+ git clone https://github.com/ohmyzsh/ohmyzsh.git /root-layer/.oh-my-zsh
+
+COPY root/ /root-layer/
+
+# runtime stage
 FROM scratch
 
 LABEL maintainer="MiguelNdeCarvalho"
 
-# copy local files
-COPY root/ /
+# Add files from buildstage
+COPY --from=buildstage /root-layer/ /
diff --git a/README.md b/README.md
index fe2dea0..a975981 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,9 @@
-# Prolog - Docker mod for code-server
+# Zsh / Oh My Zsh - Docker mod for code-server
 
-This mod adds `zsh` to code-server, to be installed/updated during container start.
+This mod adds `zsh` and `Oh My Zsh` to code-server, to be installed/updated during container start.
 
 In code-server docker arguments, set an environment variable `DOCKER_MODS=linuxserver/mods:code-server-zsh`
 
 If adding multiple mods, enter them in an array separated by `|`, such as `DOCKER_MODS=linuxserver/mods:code-server-zsh|linuxserver/mods:code-server-prolog`
+
+`Oh My Zsh` is installed at `/config/.oh-my-zsh` and can be updated via `omz update` in terminal.
diff --git a/root/etc/cont-init.d/98-zsh b/root/etc/cont-init.d/98-zsh
index 4479221..284050e 100644
--- a/root/etc/cont-init.d/98-zsh
+++ b/root/etc/cont-init.d/98-zsh
@@ -1,12 +1,21 @@
 #!/usr/bin/with-contenv bash
 
-# First install ZSH
+# Install ZSH
+echo "**** Installing zsh ****"
 apt-get update && apt-get install -y \
   zsh
 
-# Install oh-my-zsh
-sudo -u abc "git clone https://github.com/ohmyzsh/ohmyzsh.git ~/.oh-my-zsh"
-sudo -u abc "cp ~/.oh-my-zsh/templates/zshrc.zsh-template ~/.zshrc"
+# Install oh-my-zsh if not installed already
+if [ -d /config/.oh-my-zsh ]; then
+  echo "**** Oh My Zsh is already installed. You can update it via \"omz update\" in terminal. ****"
+elif [ -d /.oh-my-zsh ]; then
+  echo "**** Installing Oh My Zsh. You can update it via \"omz update\" in terminal. ****"
+  mv /.oh-my-zsh /config/.oh-my-zsh
+  cp /config/.oh-my-zsh/templates/zshrc.zsh-template /config/.zshrc
+else
+  echo "**** It looks like Oh My Zsh was manually deleted after install. Please recreate the code-server container to reinstall Oh My Zsh. ****"
+fi
 
-# Setup shell as default
-chsh -s $(which zsh) abc
+chown -R abc:abc \
+  /config/.oh-my-zsh \
+  /config/.zshrc