# Use the official ROCm base image
FROM ubuntu:26.04 AS reqs

# Set build variables
ENV PIP_BREAK_SYSTEM_PACKAGES=1
ENV PYTHONUNBUFFERED=1
ENV DEBIAN_FRONTEND=noninteractive

# Set path variables
ENV ROCM_PATH=/usr
ENV HIP_PATH=/usr

# Install libraries
RUN apt update -y && apt install -y --no-install-recommends \
    rocm \
    python3-pip \
    python3-dev \
    python3-torch-rocm \
    libhipblas-dev \
    librocblas-dev \
    git \
    libgl1 \
    libglib2.0-bin \
    cmake \
    ninja-build \
    pkg-config \
    libavcodec-dev \
    libavformat-dev \
    libswresample-dev \
    build-essential \
    g++ \
    ninja-build \
    && rm -rf /var/lib/apt/lists/*

# SymLink for stability
RUN ln -s /usr /opt/rocm

FROM reqs AS torch-audio

WORKDIR /builder

RUN git clone --depth 1 --branch v2.9.1 https://github.com/pytorch/audio.git && \
    cd audio && \
    git submodule update --init --recursive && \
    # Set path variables
    export CMAKE_PREFIX_PATH=/usr/lib/python3/dist-packages/torch/share/cmake && \
    export Torch_DIR=/usr/lib/python3/dist-packages/torch/share/cmake/Torch && \
    export hip_DIR=/usr/lib/aarch64-linux-gnu/cmake/hip && \
    export CMAKE_PREFIX_PATH="/usr/lib/python3/dist-packages/torch/share/cmake:/usr/lib/aarch64-linux-gnu/cmake/hip" && \
    # Set environment variables
    export PYTORCH_ROCM_ARCH=gfx1100 && \
    export HIP_VISIBLE_DEVICES=0 && \
    # Execute build
    USE_ROCM=1 BUILD_ROCM=1 CC=gcc CXX=g++ python3 setup.py install && \
    cd .. && rm -rf audio

FROM torch-audio AS runner

WORKDIR /app

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

# Create sub dependancy
RUN grep -vE 'torch|nvidia|cuda|triton|kornia|spandrel' requirements.txt > req_no_torch.txt

# Install non torch dependant requirements
RUN pip3 install --no-cache-dir --break-system-packages -r req_no_torch.txt

# Manually install the sensitive ones with --no-deps
RUN pip3 install --no-deps --break-system-packages torchsde kornia spandrel

# 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"]