{megapaks} Update PyTorch 2.9.0 to 2.9.1

Author: YanWenKun

cu126-megapak/Dockerfile

@@ -235,7 +235,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     && pip install \
         --upgrade pip wheel setuptools \
     && pip install \
-        --dry-run xformers==0.0.33.post1 torchvision torchaudio \
+        --dry-run xformers==0.0.33.post2 torchvision torchaudio \
         --index-url https://download.pytorch.org/whl/cu126
 
 # Here's some hack. To reduce image size & single layer size, we:
@@ -246,14 +246,14 @@ RUN --mount=type=cache,target=/root/.cache/pip \
 # Note: xFormers is not always released with latest PyTorch, so using `pip install --no-deps` is not a good idea.
 RUN --mount=type=cache,target=/root/.cache/pip \
     pip install \
-        xformers==0.0.33.post1 torchvision torchaudio \
+        xformers==0.0.33.post2 torchvision torchaudio \
         --index-url https://download.pytorch.org/whl/cu126 \
 && pip uninstall --yes torch \
 && find /usr/local/lib/python3.12/site-packages/nvidia/ -mindepth 1 -maxdepth 1 ! -name 'nccl' ! -name 'nvshmem' -exec rm -rfv {} +
 
 RUN --mount=type=cache,target=/root/.cache/pip \
     pip install \
-        xformers==0.0.33.post1 torchvision torchaudio \
+        xformers==0.0.33.post2 torchvision torchaudio \
         --index-url https://download.pytorch.org/whl/cu126
 
 ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}\
@@ -313,7 +313,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
 https://github.com/nunchaku-tech/nunchaku/releases/download/v1.0.2/nunchaku-1.0.2+torch2.9-cp312-cp312-linux_x86_64.whl \
     # FlashAttention (version pair with xFormers, binary pair with PyTorch & CUDA)
     && pip install \
-https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.4.17/flash_attn-2.8.3+cu126torch2.9-cp312-cp312-linux_x86_64.whl
+https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.5.4/flash_attn-2.8.3+cu126torch2.9-cp312-cp312-linux_x86_64.whl
 
 ################################################################################
 # Bundle ComfyUI in the image

cu128-megapak/Dockerfile

@@ -235,7 +235,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     && pip install \
         --upgrade pip wheel setuptools \
     && pip install \
-        --dry-run xformers==0.0.33.post1 torchvision torchaudio \
+        --dry-run xformers==0.0.33.post2 torchvision torchaudio \
         --index-url https://download.pytorch.org/whl/cu128
 
 # Here's some hack. To reduce image size & single layer size, we:
@@ -246,14 +246,14 @@ RUN --mount=type=cache,target=/root/.cache/pip \
 # Note: xFormers is not always released with latest PyTorch, so using `pip install --no-deps` is not a good idea.
 RUN --mount=type=cache,target=/root/.cache/pip \
     pip install \
-        xformers==0.0.33.post1 torchvision torchaudio \
+        xformers==0.0.33.post2 torchvision torchaudio \
         --index-url https://download.pytorch.org/whl/cu128 \
 && pip uninstall --yes torch \
 && find /usr/local/lib/python3.12/site-packages/nvidia/ -mindepth 1 -maxdepth 1 ! -name 'nccl' ! -name 'nvshmem' -exec rm -rfv {} +
 
 RUN --mount=type=cache,target=/root/.cache/pip \
     pip install \
-        xformers==0.0.33.post1 torchvision torchaudio \
+        xformers==0.0.33.post2 torchvision torchaudio \
         --index-url https://download.pytorch.org/whl/cu128
 
 ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}\
@@ -313,7 +313,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
 https://github.com/nunchaku-tech/nunchaku/releases/download/v1.0.2/nunchaku-1.0.2+torch2.9-cp312-cp312-linux_x86_64.whl \
     # FlashAttention (version pair with xFormers, binary pair with PyTorch & CUDA)
     && pip install \
-https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.4.17/flash_attn-2.8.3+cu128torch2.9-cp312-cp312-linux_x86_64.whl \
+https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.5.4/flash_attn-2.8.3+cu128torch2.9-cp312-cp312-linux_x86_64.whl \
     # SageAttention
     && pip install \
 https://github.com/YanWenKun/ComfyUI-Containerfiles/releases/download/sageattn2/sageattention-2.2.0-cp312-cp312-linux_x86_64.whl \