diff --git a/Dockerfile b/Dockerfile
index 86438ae..7364996 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,17 +1,5 @@
 # Base image
-# The following docker base image is recommended by VLLM: 
-FROM runpod/base:0.4.1-cuda11.8.0
-
-# Use bash shell with pipefail option
-SHELL ["/bin/bash", "-o", "pipefail", "-c"]
-
-# Set the working directory
-WORKDIR /
-
-# Update and upgrade the system packages (Worker Template)
-ARG DEBIAN_FRONTEND=noninteractive
-
-RUN pip install torch==2.0.1 -f https://download.pytorch.org/whl/cu118
+FROM runpod/base:0.4.2-cuda11.8.0
 
 # Install Python dependencies (Worker Template)
 COPY builder/requirements.txt /requirements.txt
@@ -21,14 +9,10 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     rm /requirements.txt
 
 # Add src files (Worker Template)
-ADD src .  
-
-# Quick temporary updates
-RUN pip install git+https://github.com/runpod/runpod-python@a1#egg=runpod --compile
+ADD src .
 
 # Prepare the models inside the docker image
-ARG HUGGING_FACE_HUB_TOKEN=
-ENV HUGGING_FACE_HUB_TOKEN=$HUGGING_FACE_HUB_TOKEN
+ARG HUGGING_FACE_HUB_TOKEN
 
 # Prepare argument for the model and tokenizer
 ARG MODEL_NAME=""
@@ -61,7 +45,7 @@ ENV MODEL_NAME=$MODEL_NAME \
 ENTRYPOINT ["/entrypoint.sh"]
 
 # Run the Python script to download the model
-RUN python -u /download_model.py
+RUN python -u /download_model.py --model_name $MODEL_NAME --model_revision $MODEL_REVISION --model_base_path $MODEL_BASE_PATH --hugging_face_hub_token $HUGGING_FACE_HUB_TOKEN
 
 # Start the handler
-CMD STREAMING=$STREAMING MODEL_NAME=$MODEL_NAME MODEL_BASE_PATH=$MODEL_BASE_PATH TOKENIZER=$TOKENIZER QUANTIZATION=$QUANTIZATION python -u /handler.py 
+CMD STREAMING=$STREAMING MODEL_NAME=$MODEL_NAME MODEL_BASE_PATH=$MODEL_BASE_PATH TOKENIZER=$TOKENIZER QUANTIZATION=$QUANTIZATION python -u /handler.py
diff --git a/builder/requirements.txt b/builder/requirements.txt
index 70aa770..433ab30 100644
--- a/builder/requirements.txt
+++ b/builder/requirements.txt
@@ -1,6 +1,5 @@
-# Required Python packages get listed here, one per line.
-# Recomended to lock the version number to avoid unexpected changes.
-fastapi==0.99.1
-vllm==0.2.1.post1
-huggingface-hub==0.17.3
-runpod @ git+https://github.com/runpod/runpod-python@main#egg=runpod 
+git+https://github.com/runpod/runpod-python.git@concurrency-optimization
+torch==2.0.1 -f https://download.pytorch.org/whl/cu118
+
+vllm==0.2.2
+huggingface-hub==0.19.4
diff --git a/src/download_model.py b/src/download_model.py
index 9eb5299..fbe0a4e 100644
--- a/src/download_model.py
+++ b/src/download_model.py
@@ -1,21 +1,25 @@
 import os
+import argparse
 from huggingface_hub import snapshot_download
 
-# Get the hugging face token
-HUGGING_FACE_HUB_TOKEN = os.environ.get('HUGGING_FACE_HUB_TOKEN', None)
+
 MODEL_NAME = os.environ.get('MODEL_NAME')
 MODEL_REVISION = os.environ.get('MODEL_REVISION', "main")
 MODEL_BASE_PATH = os.environ.get('MODEL_BASE_PATH', '/runpod-volume/')
 
-# Download the model from hugging face
-download_kwargs = {}
 
-if HUGGING_FACE_HUB_TOKEN:
-    download_kwargs["token"] = HUGGING_FACE_HUB_TOKEN
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model_name', type=str, default=MODEL_NAME)
+    parser.add_argument('--model_revision', type=str, default=MODEL_REVISION)
+    parser.add_argument('--model_base_path', type=str, default=MODEL_BASE_PATH)
+    parser.add_argument('--hugging_face_hub_token', type=str, default=HUGGING_FACE_HUB_TOKEN)
 
-snapshot_download(
-    MODEL_NAME,
-    revision=MODEL_REVISION,
-    local_dir=f"{MODEL_BASE_PATH}{MODEL_NAME.split('/')[1]}",
-    **download_kwargs
-)
+    args = parser.parse_args()
+
+    snapshot_download(
+        args.model_name,
+        revision=args.model_revision,
+        local_dir=f"{args.model_base_path}{args.model_name.split('/')[1]}",
+        token=args.hugging_face_hub_token
+    )