diff --git a/comfyui-rocm-dockerfile b/comfyui-rocm-dockerfile
index e793058..904b583 100644
--- a/comfyui-rocm-dockerfile
+++ b/comfyui-rocm-dockerfile
@@ -1,6 +1,7 @@
-FROM ubuntu:24.04 AS mesa-builder
+FROM ubuntu:24.04 AS builder
 ENV DEBIAN_FRONTEND=noninteractive
 
+#install build deps
 RUN apt-get update && apt-get install -y \
     git \
     python3-mako \
@@ -33,36 +34,25 @@ RUN apt-get update && apt-get install -y \
     libx11-dev \
     libxcb-shm0-dev \
     libxrandr-dev \
-    git \
-    python3-mako \
-    python3-yaml \
-    python3-packaging \
     python3-pip \
-    ninja-build \
-    build-essential \
-    pkg-config \
-    flex \
-    bison \
-    libpciaccess-dev \
-    libexpat1-dev \
-    libelf-dev \
-    libvulkan-dev \
-    glslang-tools \
-    spirv-tools \
-    llvm-dev \
-    libclang-dev \
-    cmake \
+    curl \
+    wget \
+    golang-go \
+    ca-certificates \
+    vulkan-tools \
+    spirv-headers \
+    libshaderc-dev \
+    glslc \
     && rm -rf /var/lib/apt/lists/*
-
 RUN pip3 install --no-cache-dir --break-system-packages --upgrade meson
 
-# Build required newer libdrm from source
+# Build libdrm
 RUN git clone --depth 1 https://gitlab.freedesktop.org/mesa/drm.git /drm
 WORKDIR /drm
 RUN meson setup build -Dbuildtype=release -Dprefix=/opt/mesa
 RUN ninja -C build -j16 && ninja -C build install
 
-# Build matching LLVMSPIRVLib (SPIRV-LLVM-Translator) for LLVM 18
+# Build LLVMSPIRVLib
 RUN git clone --branch llvm_release_180 --depth 1 https://github.com/KhronosGroup/SPIRV-LLVM-Translator.git /spirv-llvm-translator
 WORKDIR /spirv-llvm-translator
 RUN cmake -B build \
@@ -72,54 +62,27 @@ RUN cmake -B build \
 RUN cmake --build build --parallel 16
 RUN cmake --install build
 
-# Build Mesa-git pointing to the custom libdrm and LLVMSPIRVLib paths
+# Build Mesa-git
 RUN git clone --depth 1 https://gitlab.freedesktop.org/mesa/mesa.git /mesa
 WORKDIR /mesa
-
 ENV PKG_CONFIG_PATH=/opt/mesa/lib/pkgconfig:/opt/mesa/lib/aarch64-linux-gnu/pkgconfig:$PKG_CONFIG_PATH
 ENV CMAKE_PREFIX_PATH=/opt/mesa:$CMAKE_PREFIX_PATH
-
 RUN meson setup build \
     -Dbuildtype=release \
     -Dgallium-drivers= \
     -Dvulkan-drivers=amd \
     -Dplatforms=x11,wayland \
     -Dprefix=/opt/mesa
-
 RUN ninja -C build --verbose
 RUN ninja -C build install
 
-# Stage 2: Build Ollama with Vulkan support
-FROM ubuntu:24.04 AS ollama-builder
-ENV DEBIAN_FRONTEND=noninteractive
-
-RUN apt-get update && apt-get install -y \
-    build-essential \
-    git \
-    curl \
-    wget \
-    cmake \
-    golang-go \
-    ca-certificates \
-    ninja-build \
-    libvulkan-dev \
-    vulkan-tools \
-    glslang-tools \
-    spirv-headers \
-    spirv-tools \
-    libshaderc-dev \
-    glslc \
-    && rm -rf /var/lib/apt/lists/*
-
+# Build Ollama
+RUN git clone --recursive https://github.com/ollama/ollama.git /app
 WORKDIR /app
-RUN git clone --recursive https://github.com/ollama/ollama.git .
-
-# Configure and compile Ollama using the Vulkan backend
 RUN cmake -B build . \
     -DOLLAMA_LLAMA_BACKENDS="vulkan" \
     -DGGML_VULKAN=ON \
     -DCMAKE_BUILD_TYPE=Release
-
 RUN cmake --build build --parallel 16
 RUN go build -o ollama .
 
@@ -138,11 +101,11 @@ RUN apt-get update && apt-get install -y \
     && rm -rf /var/lib/apt/lists/*
 
 # Copy compiled Mesa-git libraries and ICD files from Stage 1
-COPY --from=mesa-builder /opt/mesa /opt/mesa
+COPY --from=builder /opt/mesa /opt/mesa
 
 # Copy compiled Ollama binary from Stage 2
-COPY --from=ollama-builder /app/ollama /app/ollama
-COPY --from=ollama-builder /app/build /app/build
+COPY --from=builder /app/ollama /app/ollama
+COPY --from=builder /app/build /app/build
 
 EXPOSE 11434