FROM rocm/dev-ubuntu-24.04:latest

# Install system dependencies
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        ca-certificates \
        git \
        libgl1 \
        libglib2.0-0 \
        libglx-mesa0 \
        && rm -rf /var/lib/apt/lists/*

# Create User Permissions
ARG DOCKER_USER=default_user
RUN useradd -m "$DOCKER_USER"

# Clone vladmandic's stable diffusion
RUN \
  git clone --depth 1 --depth 1 https://github.com/comfyanonymous/ComfyUI

WORKDIR /ComfyUI/custom_nodes/
RUN git clone --depth 1 https://github.com/ltdrdata/ComfyUI-Manager /ComfyUI/custom_nodes/comfyui-manager && \
    git clone --depth 1 https://github.com/city96/ComfyUI-GGUF /ComfyUI/custom_nodes/ComfyUI-GGUF && \
    git clone --depth 1 https://github.com/SeanScripts/ComfyUI-Unload-Model.git /ComfyUI/custom_nodes/ComfyUI-Unload-Model && \
    git clone --depth 1 https://github.com/ssitu/ComfyUI_UltimateSDUpscale --recursive /ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale && \
    git clone --depth 1 https://github.com/ltdrdata/ComfyUI-Impact-Pack /ComfyUI/custom_nodes/ComfyUI-Impact-Pack && \
    git clone --depth 1 https://github.com/ltdrdata/ComfyUI-Impact-Subpack /ComfyUI/custom_nodes/ComfyUI-Impact-Subpack && \
    git clone --depth 1 https://github.com/talesofai/comfyui-browser /ComfyUI/custom_nodes/comfyui-browser && \
    git clone --depth 1 https://github.com/SEkINVR/ComfyUI-SaveAs.git /ComfyUI/custom_nodes/ComfyUI-SaveAs

WORKDIR /ComfyUI

RUN pip install --no-cache-dir --break-system-packages --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm6.4
#RUN pip install --no-cache-dir --break-system-packages torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.3

# Install Python dependencies
RUN pip install --no-cache-dir --break-system-packages -r requirements.txt && \
    pip install --no-cache-dir --break-system-packages -r /ComfyUI/custom_nodes/comfyui-manager/requirements.txt && \
    pip install --no-cache-dir --break-system-packages -r /ComfyUI/custom_nodes/ComfyUI-GGUF/requirements.txt  && \
    pip install --no-cache-dir --break-system-packages -r /ComfyUI/custom_nodes/ComfyUI-Impact-Pack/requirements.txt && \
    pip install --no-cache-dir --break-system-packages -r /ComfyUI/custom_nodes/ComfyUI-Impact-Subpack/requirements.txt && \
    pip install --no-cache-dir --break-system-packages -r /ComfyUI/custom_nodes/comfyui-browser/requirements.txt && \
    pip install --no-cache-dir --break-system-packages -r /ComfyUI/custom_nodes/ComfyUI-SaveAs/requirements.txt && \
    pip install --no-cache-dir --break-system-packages gitpython psutil

# Set ownership and move models directory
RUN chown -R $DOCKER_USER:$DOCKER_USER /ComfyUI && \
    mv /ComfyUI/models /ComfyUI/models.orig

# Copy start script
COPY ./start.sh /usr/local/bin/
RUN chmod +x /usr/local/bin/start.sh

# Switch to non-root user
USER $DOCKER_USER
WORKDIR /ComfyUI

CMD ["/usr/local/bin/start.sh"]