diff --git a/comfyui-rocm-dockerfile b/comfyui-rocm-dockerfile
index ae21752..3d708ac 100644
--- a/comfyui-rocm-dockerfile
+++ b/comfyui-rocm-dockerfile
@@ -1,16 +1,18 @@
 # Use the official ROCm base image
-FROM ubuntu:24.04
+FROM ubuntu:26.04
 
 # Set environment variables for 7900 XTX (gfx1100)
-ENV HSA_OVERRIDE_GFX_VERSION=11.0.0
-ENV TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
-ENV LD_LIBRARY_PATH=/usr/lib/aarch64-linux-gnu:$LD_LIBRARY_PATH
 ENV PIP_BREAK_SYSTEM_PACKAGES=1
 ENV PYTHONUNBUFFERED=1
 
 RUN apt update -y && apt install -y --no-install-recommends \
+    rocm-dev \
+    rocm-libs \
+    rocminfo \
     python3-pip \
     python3-dev \
+    python3-torch-rocm \
+    python3-torchaudio \
     git \
     libgl1 \
     libglib2.0-0 \
@@ -22,12 +24,9 @@ WORKDIR /app
 # Clone ComfyUI
 RUN git clone https://github.com/comfyanonymous/ComfyUI.git .
 
-# Install PyTorch with ROCm 6.2 support (matching the base image)
-RUN pip3 install --no-cache-dir torch torchvision torchaudio \
-    --index-url https://download.pytorch.org/whl/rocm6.2
-
 # Install ComfyUI dependencies
-RUN pip3 install --no-cache-dir -r requirements.txt
+RUN grep -vE 'torch|torchaudio|torchvision' requirements.txt > req_no_torch.txt \
+    && pip3 install --no-cache-dir -r req_no_torch.txt
 
 # Set custom Nodes
 WORKDIR /app/custom_nodes
diff --git a/rocm-docker-compose.yml b/rocm-docker-compose.yml
index 5862bb8..fd62f9b 100644
--- a/rocm-docker-compose.yml
+++ b/rocm-docker-compose.yml
@@ -11,8 +11,8 @@ services:
       - OLLAMA_KEEP_ALIVE=24h
       - HSA_OVERRIDE_GFX_VERSION=11.0.0 # This is for the RX 7900XTX
     devices:
-      - /dev/kfd:/dev/kfd # Shared surface driver for ROCm
-      - /dev/dri:/dev/dri # Direct Rendering Infrastructure for GPU access
+      - /dev/kfd:/dev/kfd
+      - /dev/dri:/dev/dri
 
   ollama-webui:
     container_name: ai-ollama-webui