install flash attention

Dockerfile

@@ -45,6 +45,11 @@ COPY --chown=user requirements.txt .
 RUN --mount=type=cache,target=$UV_CACHE_DIR,uid=1000,gid=1000 \
     /uv pip install -r requirements.txt
 
+# Conditionally install flash-attn if CUDA is available
+RUN --mount=type=cache,target=$UV_CACHE_DIR,uid=1000,gid=1000 \
+    python -c "import torch; exit(0 if torch.cuda.is_available() else 1)" && \
+    /uv pip install flash-attn --no-build-isolation || echo "CUDA not available, skipping flash-attn installation"
+
 # Copy application code
 COPY --chown=user . .
 

main.py

@@ -43,14 +43,6 @@ LANGUAGE = os.getenv("LANGUAGE", "english").lower()
 
 device = get_device(force_cpu=False)
 
-# Install Flash Attention 2 if device is "cuda"
-if device == "cuda":
-    subprocess.run(
-        ["pip", "install", "flash-attn", "--no-build-isolation"],
-        env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"},
-        shell=True,
-    )
-
 torch_dtype, np_dtype = get_torch_and_np_dtypes(device, use_bfloat16=False)
 logger.info(f"Using device: {device}, torch_dtype: {torch_dtype}, np_dtype: {np_dtype}")