Update requirements.txt

requirements.txt

@@ -1,3 +1,4 @@
+# Core dependencies
 anyio==4.8.0
 fastapi==0.116.1
 gradio==5.16.0
@@ -6,16 +7,29 @@ huggingface-hub==0.28.1
 numpy==2.2.6
 safetensors==0.5.3
 soundfile==0.13.1
-torch==2.6.0
-torchaudio==2.6.0
-triton==3.2.0
-transformers==4.49.0
 typing-extensions==4.12.2
 uvicorn==0.35.0
+transformers==4.49.0
 spaces
-# Install mamba-ssm from wheel
-https://github.com/state-spaces/mamba/releases/download/v2.2.4/mamba_ssm-2.2.4+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
-# Install causal-conv1d from wheel  
+
+# PyTorch and related
+torch==2.6.0
+torchaudio==2.6.0
+torchvision==0.21.0
+
+# IMPORTANT: Install flash-attn BEFORE mamba-ssm
+# flash-attn must be installed from wheel for speed
+https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.2.post1/flash_attn-2.7.2.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
+
+# Then install causal-conv1d
 https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.5.0.post8/causal_conv1d-1.5.0.post8+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
-# Install zonos if it's available
-zonos @ git+https://github.com/Zyphra/Zonos.git
+
+# Finally install mamba-ssm (it needs flash-attn to be already installed)
+https://github.com/state-spaces/mamba/releases/download/v2.2.4/mamba_ssm-2.2.4+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
+
+# Zonos package
+zonos @ git+https://github.com/Zyphra/Zonos.git
+
+# Additional dependencies that might be needed
+einops
+ninja