diff --git a/comfyui-rocm-dockerfile b/comfyui-rocm-dockerfile
index 677d439..a5e540b 100644
--- a/comfyui-rocm-dockerfile
+++ b/comfyui-rocm-dockerfile
@@ -1,4 +1,4 @@
-# Stage 1: Build Ollama with Vulkan support for ARM64
+# Stage 1: Build Ollama with Vulkan support for ARM64/AMD64
 FROM ubuntu:24.04 AS builder
 
 ENV DEBIAN_FRONTEND=noninteractive
@@ -16,10 +16,12 @@ RUN apt-get update && apt-get install -y \
 WORKDIR /app
 RUN git clone --recursive https://github.com/ollama/ollama.git .
 
-# Build Ollama with Vulkan support enabled
-ENV OLLAMA_SET_VULKAN=1
-RUN go generate ./... && \
-    go build -tags vulkan .
+# Configure and build using CMake for Vulkan backend support
+RUN cmake -B build -D CMAKE_BUILD_TYPE=Release -D OLLAMA_LLAMA_BACKENDS="vulkan"
+RUN cmake --build build --parallel $(nproc)
+
+# Build the final Go binary
+RUN go build -o ollama .
 
 # Stage 2: Minimal Runtime Image
 FROM ubuntu:24.04
@@ -33,9 +35,11 @@ RUN apt-get update && apt-get install -y \
     ca-certificates \
     && rm -rf /var/lib/apt/lists/*
 
-# Copy compiled binaries from builder stage
+# Copy the compiled binary from the builder stage
 COPY --from=builder /app/ollama /usr/local/bin/ollama
-COPY --from=builder /app/dist /usr/local/bin/dist
+
+# Copy built backend libraries if generated under build output
+COPY --from=builder /app/build/lib/ollama /usr/local/lib/ollama
 
 EXPOSE 11434
 ENV OLLAMA_HOST=0.0.0.0:11434