diff --git a/Dockerfile.nix b/Dockerfile.nix
index 29158bea..2e85dd00 100644
--- a/Dockerfile.nix
+++ b/Dockerfile.nix
@@ -1,8 +1,7 @@
 # Build the image and get out the docker file:
 #
-# docker build -t tgi-nix -f Dockerfile.nix
-# docker run --rm --volume $PWD/data:/data tgi-nix cp -H result /data/tgi-docker.tar.gz
-# docker load < tgi-docker.tar.gz
+# docker build -t tgi-nix-builder -f Dockerfile.nix
+# docker run tgi-nix-builder | docker load
 
 FROM nixos/nix:2.18.8
 RUN echo "experimental-features = nix-command flakes" >> /etc/nix/nix.conf
@@ -10,4 +9,5 @@ RUN nix profile install nixpkgs#cachix
 RUN cachix use text-generation-inference
 WORKDIR /root
 ADD . .
-RUN nix build .#dockerImage
+RUN nix build .#dockerImageStreamed
+ENTRYPOINT ./result
diff --git a/flake.nix b/flake.nix
index ddd31a18..22fa8cfc 100644
--- a/flake.nix
+++ b/flake.nix
@@ -153,18 +153,13 @@
             '';
           };
 
-          dockerImage = pkgs.dockerTools.buildLayeredImage {
-            name = "tgi-docker";
-            tag = "latest";
-            config = {
-              EntryPoint = [ "${default}/bin/text-generation-inference" ];
-              Env = [
-                "HF_HOME=/data"
-                "PORT=80"
-              ];
+          dockerImage = pkgs.callPackage nix/docker.nix {
+            text-generation-inference = default;
+          };
 
-            };
-            contents = [ pkgs.cacert ];
+          dockerImageStreamed = pkgs.callPackage nix/docker.nix {
+            text-generation-inference = default;
+            stream = true;
           };
         };
       }
diff --git a/nix/docker.nix b/nix/docker.nix
new file mode 100644
index 00000000..c4b1d899
--- /dev/null
+++ b/nix/docker.nix
@@ -0,0 +1,23 @@
+{
+  dockerTools,
+  cacert,
+  text-generation-inference,
+  stream ? false,
+}:
+
+let
+  build = if stream then dockerTools.streamLayeredImage else dockerTools.buildLayeredImage;
+in
+build {
+  name = "tgi-docker";
+  tag = "latest";
+  config = {
+    EntryPoint = [ "${text-generation-inference}/bin/text-generation-inference" ];
+    Env = [
+      "HF_HOME=/data"
+      "PORT=80"
+    ];
+
+  };
+  contents = [ cacert ];
+}