# ComfyUI with ROCm support for AMD Ryzen AI Max+ 395
# Based on official AMD ROCm blog guide: [URL REMOVED]
# Optimized for layer caching - changes are ordered from least to most frequent

FROM rocm/pytorch:rocm7.1_ubuntu24.04_py3.12_pytorch_release_2.9.1

# Set environment variable (matching official AMD guide)
ENV COMFYUI_PATH=/workload/ComfyUI

# Layer 1: Upgrade pip and install base tools (rarely changes)
RUN pip install --upgrade pip wheel setuptools

# Layer 2: Install ROCm PyTorch wheels from AMD repository (ROCm 7.2.0)
# This is a large download but rarely changes - cache separately
RUN cd /tmp && \
    wget -q "[URL REMOVED] && \
    wget -q "[URL REMOVED] && \
    wget -q "[URL REMOVED] && \
    wget -q "[URL REMOVED] && \
    pip uninstall -y torch torchvision torchaudio triton || true && \
    pip install --break-system-packages --no-deps /tmp/triton-3.5.1+rocm7.2.0.gita272dfa8-cp312-cp312-linux_x86_64.whl && \
    pip install --break-system-packages --no-deps /tmp/torch-2.9.1+rocm7.2.0.lw.git7e1940d4-cp312-cp312-linux_x86_64.whl && \
    pip install --break-system-packages --no-deps /tmp/torchvision-0.24.0+rocm7.2.0.gitb919bd0c-cp312-cp312-linux_x86_64.whl && \
    pip install --break-system-packages --no-deps /tmp/torchaudio-2.9.0+rocm7.2.0.gite3c6ee2b-cp312-cp312-linux_x86_64.whl && \
    rm -f /tmp/*.whl

# Layer 3: Install common Python dependencies (changes less frequently than ComfyUI code)
RUN pip install --break-system-packages matplotlib pandas simpleeval

# Layer 4: Clone ComfyUI and custom nodes (changes when ComfyUI updates)
# Use --depth 1 to get latest version (ComfyUI uses master branch, not main)
RUN git clone --depth 1 [URL REMOVED] $COMFYUI_PATH && \
    git clone --depth 1 [URL REMOVED] $COMFYUI_PATH/custom_nodes/comfyui-manager && \
    git clone --depth 1 [URL REMOVED] $COMFYUI_PATH/custom_nodes/rgthree-comfy

# Set ComfyUI as working directory
WORKDIR $COMFYUI_PATH

# Layer 5: Install ComfyUI requirements (changes when requirements.txt updates)
RUN pip install --break-system-packages -r $COMFYUI_PATH/requirements.txt && \
    pip install --break-system-packages -r $COMFYUI_PATH/custom_nodes/comfyui-manager/requirements.txt

# Layer 6: Copy scripts (changes when scripts are modified)
# Copy entrypoint first (changes less frequently)
# Updated: 2026-01-25 - Added memory optimizations (reserve-vram 16, memory_usage_factor 0.2)
COPY entrypoint.sh /workspace/entrypoint.sh
RUN chmod +x /workspace/entrypoint.sh

# Copy other scripts
COPY download-flux1.py /workspace/download-flux1.py
COPY auto-load-workflow.py /workspace/auto-load-workflow.py
RUN chmod +x /workspace/auto-load-workflow.py

# Expose ComfyUI port
EXPOSE 8188

# Use custom entrypoint (official guide uses /bin/bash, but we use our entrypoint)
CMD ["/workspace/entrypoint.sh"]