diff --git a/.github/workflows/build-and-push.yml b/.github/workflows/build-and-push.yml
index ea6ca68..4cc89bc 100644
--- a/.github/workflows/build-and-push.yml
+++ b/.github/workflows/build-and-push.yml
@@ -38,7 +38,7 @@ jobs:
           password: ${{ secrets.DOCKERHUB_TOKEN }}
       - uses: docker/build-push-action@v6
         with:
-          context: ./wrapper
+          context: ./.runpod
           build-args: OLLAMA_VERSION=${{ inputs.ollama_version }}
           push: true
           tags: svenbrnn/runpod-ollama:${{ inputs.ollama_version }}${{ inputs.next_version_suffix }},svenbrnn/runpod-ollama:latest
diff --git a/wrapper/Dockerfile b/.runpod/Dockerfile
similarity index 90%
rename from wrapper/Dockerfile
rename to .runpod/Dockerfile
index fa55087..1687f9e 100644
--- a/wrapper/Dockerfile
+++ b/.runpod/Dockerfile
@@ -32,7 +32,7 @@ RUN apt-get update --yes --quiet && DEBIAN_FRONTEND=noninteractive apt-get insta
 WORKDIR /work
 
 # Add my src as /work
-ADD src /work
+ADD engine.py handler.py requirements.txt start.sh test_input.json utils.py /work
 
 # Set defaut ollama models directory to /runpod-volume where runpod will mount the volume by default
 ENV OLLAMA_MODELS="/runpod-volume"
diff --git a/wrapper/src/engine.py b/.runpod/engine.py
similarity index 100%
rename from wrapper/src/engine.py
rename to .runpod/engine.py
diff --git a/wrapper/src/runpod_wrapper.py b/.runpod/handler.py
similarity index 100%
rename from wrapper/src/runpod_wrapper.py
rename to .runpod/handler.py
diff --git a/wrapper/src/requirements.txt b/.runpod/requirements.txt
similarity index 100%
rename from wrapper/src/requirements.txt
rename to .runpod/requirements.txt
diff --git a/wrapper/src/start.sh b/.runpod/start.sh
similarity index 96%
rename from wrapper/src/start.sh
rename to .runpod/start.sh
index 1ce9fd7..c557a56 100644
--- a/wrapper/src/start.sh
+++ b/.runpod/start.sh
@@ -37,4 +37,4 @@ else
     ollama pull $MODEL_NAME
 fi
 
-python -u runpod_wrapper.py $1
\ No newline at end of file
+python -u handler.py $1
\ No newline at end of file
diff --git a/wrapper/src/test_input.json b/.runpod/test_input.json
similarity index 100%
rename from wrapper/src/test_input.json
rename to .runpod/test_input.json
diff --git a/wrapper/src/utils.py b/.runpod/utils.py
similarity index 100%
rename from wrapper/src/utils.py
rename to .runpod/utils.py
diff --git a/README.md b/README.md
index 93eecac..df4bf1c 100644
--- a/README.md
+++ b/README.md
@@ -5,6 +5,8 @@
 Start a runpod serverless with the docker container ``svenbrnn/runpod-ollama:latest``. Set ``MODEL_NAME`` environment to a model from ollama.com to automatically download a model.
 A mounted volume will be automatically used.
 
+[![RunPod](https://api.runpod.io/badge/SvenBrnn/runpod-worker-ollama)](https://www.runpod.io/console/hub/SvenBrnn/runpod-worker-ollama)
+
 ## Environment variables
 
 | Variable Name | Description                              | Default Value       |