From 8a388f3d9249a2d6aac173317ddfa4669fee7ad4 Mon Sep 17 00:00:00 2001
From: owner <owner@felidae.xyz>
Date: Mon, 17 Nov 2025 02:03:28 -0500
Subject: [PATCH] nvidia updates

---
 services/nvidia/Dockerfile | 3 +++
 services/nvidia/start.sh   | 2 +-
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/services/nvidia/Dockerfile b/services/nvidia/Dockerfile
index 2f34c17..83cad0c 100644
--- a/services/nvidia/Dockerfile
+++ b/services/nvidia/Dockerfile
@@ -40,6 +40,9 @@ RUN pip install --break-system-packages -r requirements.txt \
     && pip install --break-system-packages -r /ComfyUI/custom_nodes/ComfyUI-SaveAs/requirements.txt \
     && pip install --break-system-packages gitpython psutil
 
+# Install prebuilt flash-attn
+RUN pip install --break-system-packages https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.4.22/flash_attn-2.8.1+cu130torch2.9-cp312-cp312-linux_x86_64.whl
+
 # Set ownership and move models directory
 RUN chown -R $DOCKER_USER:$DOCKER_USER /ComfyUI && \
     mv /ComfyUI/models /ComfyUI/models.orig
diff --git a/services/nvidia/start.sh b/services/nvidia/start.sh
index 5446056..94e8557 100644
--- a/services/nvidia/start.sh
+++ b/services/nvidia/start.sh
@@ -2,4 +2,4 @@
 echo "COPYING"
 cp -rf /ComfyUI/models.orig/* /ComfyUI/models/
 echo "START"
-python3 main.py "${CLI_ARGS}" --listen
+python3 main.py --listen "${CLI_ARGS}"
\ No newline at end of file