diff --git a/Dockerfile b/Dockerfile
index 94236d3..de6befa 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -50,6 +50,9 @@ ARG MODEL_TYPE
 # Change working directory to ComfyUI
 WORKDIR /comfyui
 
+# Create necessary directories
+RUN mkdir -p models/checkpoints models/vae
+
 # Download checkpoints/vae/LoRA to include in image based on model type
 RUN if [ "$MODEL_TYPE" = "sdxl" ]; then \
       wget -O models/checkpoints/sd_xl_base_1.0.safetensors https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors && \