51 lines
1.7 KiB
Docker
51 lines
1.7 KiB
Docker
FROM rocm/pytorch:rocm7.2.4_ubuntu24.04_py3.12_pytorch_release_2.9.1
|
|
|
|
# "Installing" flash-attention
|
|
# - Not sure if branch `main_perf` is actually better,
|
|
# but everyone seems to be using it, so ¯\_(ツ)_/¯
|
|
# - As far as I understood, we're not actually installing-installing
|
|
# flash-attention, but we're telling it to use the Triton backend/implementation
|
|
# which is already included in the rocm/pytorch image (built by amd),
|
|
# so the lines below are executed pretty fast.
|
|
# - That ^ is what the env variable seems to be for.
|
|
|
|
ENV FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE
|
|
|
|
RUN cd /opt && \
|
|
git clone https://github.com/ROCm/flash-attention.git && \
|
|
cd flash-attention && \
|
|
git checkout main_perf && \
|
|
python setup.py install
|
|
|
|
# Cloning and installing ComfyUI in case the user doesn't provide their own
|
|
# - Nothing unusual here afaik
|
|
|
|
RUN cd /opt && \
|
|
git clone https://github.com/comfyanonymous/ComfyUI ComfyUI-pre-cloned && \
|
|
cd ComfyUI-pre-cloned && \
|
|
pip3 install -r requirements.txt && \
|
|
pip3 install -r manager_requirements.txt
|
|
|
|
# Some utilities to make life/debugging easier
|
|
# Feel free to remove these if you're building from scratch locally.
|
|
|
|
RUN mkdir -p /opt/comfyui-gfx1151-utils
|
|
|
|
WORKDIR /opt/comfyui-gfx1151-utils
|
|
|
|
ADD scripts/check-comfyui.sh check-comfyui.sh
|
|
RUN chmod +x check-comfyui.sh
|
|
|
|
ADD scripts/test-pytorch.sh test-pytorch.sh
|
|
RUN chmod +x test-pytorch.sh
|
|
|
|
ADD scripts/test-pytorch-flashattention.py test-pytorch-flashattention.py
|
|
RUN chmod +x test-pytorch-flashattention.py
|
|
|
|
# Run ComfyUI
|
|
# - Also nothing unusual here afaik
|
|
|
|
EXPOSE 8188
|
|
|
|
CMD /opt/comfyui-gfx1151-utils/check-comfyui.sh && python3 /opt/ComfyUI/main.py --listen 0.0.0.0 --use-flash-attention --gpu-only --enable-manager
|