diff --git a/llama_toolchain/distribution/build_image.sh b/llama_toolchain/distribution/build_image.sh
index 10463727d..5c4a7ccf3 100644
--- a/llama_toolchain/distribution/build_image.sh
+++ b/llama_toolchain/distribution/build_image.sh
@@ -1,20 +1,22 @@
 #!/bin/bash
 
-if [ "$#" -ne 4 ]; then
-  echo "Usage: $0 <image_name> <base_image> <pip_dependencies> <entrypoint_command>"
-  echo "Example: $0 my-fastapi-app python:3.9-slim 'fastapi uvicorn' 'python3 -m llama_toolchain.distribution.server --port 8000'"
+if [ "$#" -ne 3 ]; then
+  echo "Usage: $0 <image_name> <base_image> <pip_dependencies>
+  echo "Example: $0 my-fastapi-app python:3.9-slim 'fastapi uvicorn'
   exit 1
 fi
 
 IMAGE_NAME=$1
 BASE_IMAGE=$2
 PIP_DEPENDENCIES=$3
-ENTRYPOINT_COMMAND=$4
 
 set -euo pipefail
 
+PORT=8001
+
 SCRIPT_DIR=$(dirname "$(readlink -f "$0")")
-SOURCE_DIR=$(dirname $(dirname "$SCRIPT_DIR"))
+SOURCE_DIR=$(dirname $(dirname $(dirname "$SCRIPT_DIR")))
+echo $SOURCE_DIR
 
 TEMP_DIR=$(mktemp -d)
 echo "Created temporary directory: $TEMP_DIR"
@@ -22,12 +24,14 @@ echo "Created temporary directory: $TEMP_DIR"
 cat <<EOF >"$TEMP_DIR/Dockerfile"
 FROM $BASE_IMAGE
 WORKDIR /app
-COPY llama_toolchain /app
+COPY llama-stack/llama_toolchain /app/llama_toolchain
+COPY llama-models/models /app/llama_models
 
-RUN pip install --no-cache-dir $PIP_DEPENDENCIES
+RUN pip install $PIP_DEPENDENCIES
+
+EXPOSE $PORT
+ENTRYPOINT ["python3", "-m", "llama_toolchain.distribution.server", "--port", "$PORT"]
 
-EXPOSE 8000
-CMD $ENTRYPOINT_COMMAND
 EOF
 
 echo "Dockerfile created successfully in $TEMP_DIR/Dockerfile"
diff --git a/llama_toolchain/distribution/run_image.sh b/llama_toolchain/distribution/run_image.sh
new file mode 100644
index 000000000..d6e003d25
--- /dev/null
+++ b/llama_toolchain/distribution/run_image.sh
@@ -0,0 +1,11 @@
+#!/bin/bash
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+set -euo pipefail
+
+podman run -it -p 8001:8001 -v ~/.llama/test.yaml:/app/test.yaml test-image --yaml_config /app/test.yaml --disable-ipv6