# --- Stage 1: The Builder ---
FROM ubuntu:24.04 AS builder

ENV DEBIAN_FRONTEND=noninteractive
RUN apt update && apt install -y \
    git cmake ninja-build python3-pip python3-venv \
    libnuma-dev libpci-dev libelf-dev libdrm-dev \
    build-essential g++-13 gcc-13

WORKDIR /src
RUN git clone https://github.com/ROCm/TheRock.git
WORKDIR /src/TheRock

# Configure for gfx1100 (7900 XTX) on Aarch64
# We use a venv to manage the build tools
RUN python3 -m venv venv && \
    ./venv/bin/pip install --upgrade pip && \
    ./venv/bin/pip install -r requirements.txt

# Compile a subset of ROCm (HIP + KFD) needed for runtime
RUN cmake -S . -B build \
    -G Ninja \
    -DTHEROCK_AMDGPU_TARGETS="gfx1100" \
    -DCMAKE_INSTALL_PREFIX=/opt/rocm \
    -DTHEROCK_ENABLE_ALL=OFF \
    -DTHEROCK_ENABLE_HIP=ON \
    -DTHEROCK_ENABLE_ROCBLAS=ON 

RUN cmake --build build --target install

# --- Stage 2: Final Runtime ---
FROM ubuntu:24.04

# Core Environment Variables
ENV HSA_OVERRIDE_GFX_VERSION=11.0.0
ENV TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
ENV ROCM_PATH=/opt/rocm
ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/lib/aarch64-linux-gnu:$LD_LIBRARY_PATH
ENV PIP_BREAK_SYSTEM_PACKAGES=1
ENV PYTHONUNBUFFERED=1

# Copy compiled ROCm from builder
COPY --from=builder /opt/rocm /opt/rocm

RUN apt update -y && apt install -y --no-install-recommends \
    python3-pip python3-dev git libgl1 libglib2.0-0 \
    libnuma1 libpci3 libelf1 \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /app

# Clone ComfyUI
RUN git clone https://github.com/comfyanonymous/ComfyUI.git .

# Install PyTorch Aarch64 ROCm Wheels
# NOTE: As of 2026, check the PyTorch nightly/stable index for the Aarch64-specific ROCm wheels
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

# Setup Manager
WORKDIR /app/custom_nodes
RUN git clone https://github.com/Comfy-Org/ComfyUI-Manager.git

WORKDIR /app
EXPOSE 8188

CMD ["python3", "main.py", "--listen", "0.0.0.0"]